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


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

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

    // 在這個(gè)方法里面計(jì)算好各個(gè)cell的LayoutAttributes 對于瀑布流布局, 只需要更改LayoutAttributes.frame即可
    // 在每次collectionView的data(init delete insert reload)變化的時(shí)候都會調(diào)用這個(gè)方法準(zhǔn)備布局
    override func prepare() {
        super.prepare()
        layoutAttributes = computeLayoutAttributes()
        // 設(shè)置為當(dāng)前屏幕的寬度
        oldScreenWidth = UIScreen.main.bounds.width //deleted ()
    }
    
    // Apple建議要重寫這個(gè)方法, 因?yàn)槟承┣闆r下(delete insert...)系統(tǒng)可能需要調(diào)用這個(gè)方法來布局
    
    override func layoutAttributesForItem(at indexPath: IndexPath) -> UICollectionViewLayoutAttributes? {
        return layoutAttributes[indexPath.row]

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