原文// github: https://github.com/jasnig
// 簡書: http://www.lxweimin.com/p/b84f4dd96d0c


protocol WaterFallLayoutDelegate: NSObjectProtocol {
    // 用來設置每一個cell的高度
    func heightForItemAtIndexPath(indexPath: IndexPath) -> CGFloat
}

class WaterFallLayout: UICollectionViewLayout {
    /// 共有多少列
    var numberOfColums = 0 {
        didSet {
            // 初始化為0
            for _ in 0..<numberOfColums {
                maxYOfColums.append(0)
            }
        }
    }
    /// cell之間的間隙 默認為5.0
    var itemSpace: CGFloat = 5.0
    
    weak var delegate: WaterFallLayoutDelegate?
    // 當item比較少(幾百個)的時候建議緩存
    // 當有成千個item的時候建議其他方式計算
    private var layoutAttributes: [UICollectionViewLayoutAttributes] = []
    // 初始都設置為0
    private var maxYOfColums: [CGFloat] = []
    /// 用于記錄之前屏幕的寬度 便于在旋轉的時候刷新視圖
    private var oldScreenWidth: CGFloat = 0.0
    

    // 在這個方法里面計算好各個cell的LayoutAttributes 對于瀑布流布局, 只需要更改LayoutAttributes.frame即可
    // 在每次collectionView的data(init delete insert reload)變化的時候都會調用這個方法準備布局
    override func prepare() {
        super.prepare()
        layoutAttributes = computeLayoutAttributes()
        // 設置為當前屏幕的寬度
        oldScreenWidth = UIScreen.main.bounds.width //deleted ()
    }
    
    // Apple建議要重寫這個方法, 因為某些情況下(delete insert...)系統可能需要調用這個方法來布局
    
    override func layoutAttributesForItem(at indexPath: IndexPath) -> UICollectionViewLayoutAttributes? {
        return layoutAttributes[indexPath.row]

    }
    
    // 必須重寫這個方法來返回計算好的LayoutAttributes
    override func layoutAttributesForElements(in rect: CGRect) -> [UICollectionViewLayoutAttributes]? {
        // 返回計算好的layoutAttributes
        return layoutAttributes
    }
    
    // 返回collectionView的ContentSize -> 滾動范圍
    //changed this -> override func collectionViewContentSize() -> CGSize to this->
    override var collectionViewContentSize:  CGSize {
        let maxY = maxYOfColums.max()!
        return CGSize(width: 0.0, height: maxY)
    }
    
    // 當collectionView的bounds(滾動, 或者frame變化)發生改變的時候就會調用這個方法
    override func shouldInvalidateLayout(forBoundsChange newBounds: CGRect) -> Bool {
        // 旋轉屏幕后刷新視圖
        return newBounds.width != oldScreenWidth
        
    }
    // 計算所有的UICollectionViewLayoutAttributes
    func computeLayoutAttributes() -> [UICollectionViewLayoutAttributes] {
        //總共的item數量
        let totalNums = collectionView!.numberOfItems(inSection: 0)
        //每一個item的寬度 = （collectionview的寬度 - cell之間的間隙 默認為5.0 ）/  列數
        let width = (collectionView!.bounds.width - itemSpace * CGFloat(numberOfColums + 1)) / CGFloat(numberOfColums)
        
        var x: CGFloat
        var y: CGFloat
        var height: CGFloat
        var currentColum: Int
        var indexPath: IndexPath
        var attributesArr: [UICollectionViewLayoutAttributes] = []
        //必須設置代理
        guard let unwapDelegate = delegate else {
            assert(false, "需要設置代理")
            return attributesArr
        }
        
        for index in 0..<numberOfColums {
            self.maxYOfColums[index] = 0
        }
        for currentIndex in 0..<totalNums {
            indexPath = IndexPath(item: currentIndex, section: 0)
            
            height = unwapDelegate.heightForItemAtIndexPath(indexPath: indexPath)
            
            if currentIndex < numberOfColums {// 第一行直接添加到當前的列
                currentColum = currentIndex
                
            } else {// 其他行添加到最短的那一列
                // 這里使用!會得到期望的值
                let minMaxY = maxYOfColums.min()!
                currentColum = maxYOfColums.index(of: minMaxY)!
            }
            //            currentColum = currentIndex % numberOfColums
            x = itemSpace + CGFloat(currentColum) * (width + itemSpace)
            // 每個cell的y
            y = itemSpace + maxYOfColums[currentColum]
            // 記錄每一列的最后一個cell的最大Y
            maxYOfColums[currentColum] = y + height
            
            let attributes = UICollectionViewLayoutAttributes(forCellWith: indexPath)
            // 設置用于瀑布流效果的attributes的frame
            attributes.frame = CGRect(x: x, y: y, width: width, height: height)
            
            attributesArr.append(attributes)
        }
        return attributesArr
        
    }
}

簡單使用

var cellCount = 40
    public lazy var cellHeight:[CGFloat] = { //changed private to public
        var arr:[CGFloat] = []
        for _ in 0..<self.cellCount {
            arr.append(CGFloat(arc4random() % 150 + 40))
        }
        return arr
    }()

 func setWaterFallLayout() {
        let layout = WaterFallLayout()
        layout.delegate = self
        layout.numberOfColums = 2
        collectionView?.collectionViewLayout = layout
        
    }
extension ViewController: WaterFallLayoutDelegate {
    func heightForItemAtIndexPath(indexPath: IndexPath) -> CGFloat {
        return cellHeight[indexPath.row]
    }
}