瀑布流布局网上有很多demo，我只是使用swift重写一遍而已，原理很简单，重写UICollectionViewLayout布局

• 重写prepare方法
• 重写layoutAttributesForItem
• 重写layoutAttributesForElements
• 重写collectionViewContentSize
  原理是用数组存放每列的高度，然后找出高度最小的那一列，创建UICollectionViewLayoutAttributes布局计算出cell的frame后更新一下最短那一列的高度，所以下一次每次cell都会加载在高度最小那一列，废话不多说上代码

//
//  LFWaterFallLayout.swift
//  PXSTrain
//
//  Created by Pro on 2020/9/17.
//  Copyright © 2020 com.liufeng.mysterFeng. All rights reserved.
//

import UIKit
let LFDefaultColunmCount:Int = 3
let LFDefaultColunmMargin:CGFloat = 10.0
let LFDefaultRowMargin:CGFloat = 10.0
let LFDefaultEdgeInsets:UIEdgeInsets = UIEdgeInsets.init(top: 10, left: 10, bottom: 10, right: 10)



 @objc protocol LFWaterFallLayoutDelegate {
    //每个item的高度
    @objc func heightForItemAtIndexPath(indexPath:Int,itemWidth:CGFloat,waterFallLayout:LFWaterFallLayout) -> CGFloat
    //多少列
    @objc optional func columnCountInWaterFallLayout(waterFallLayout:LFWaterFallLayout) -> Int
    //列间距
    @objc optional func columnMarginInWaterFallLayout(waterFallLayout:LFWaterFallLayout) -> CGFloat
    //行间距
    @objc optional func rowMarginInWaterFallLayout(waterFallLayout:LFWaterFallLayout) -> CGFloat
    // 每个item的内边距
    @objc optional func edgeInsetdInWaterFallLayout(waterFallLayout:LFWaterFallLayout) -> UIEdgeInsets

}

class LFWaterFallLayout: UICollectionViewLayout {
    weak var delegate:LFWaterFallLayoutDelegate?
    //存放所有布局属性
    var attrsArr = Array<Any>()
    //存放所有列的当前高度
    var columnHeights = [CGFloat]()
    //内容的高度
    var contentHeight:CGFloat = 0
    
    var colunmCount:Int {
        get {
            return self.delegate?.columnCountInWaterFallLayout?(waterFallLayout: self) ?? LFDefaultColunmCount
        }
    }
    
    var columnMargin:CGFloat {
        get {
            return self.delegate?.columnMarginInWaterFallLayout?(waterFallLayout: self) ?? LFDefaultColunmMargin
        }
        
    }
    var rowMargin:CGFloat {
          get {
            return self.delegate?.rowMarginInWaterFallLayout?(waterFallLayout: self) ?? LFDefaultRowMargin
          }
          
      }
    var edgeInsets:UIEdgeInsets {
        get {
            return self.delegate?.edgeInsetdInWaterFallLayout?(waterFallLayout: self) ?? LFDefaultEdgeInsets
        }
    }
    
    override func prepare() {
        super.prepare()
        contentHeight = 0
        //清楚所有计算的高度
        columnHeights.removeAll()
        //设置每一列默认高度
        for _ in 0..<LFDefaultColunmCount {
            columnHeights.append(LFDefaultEdgeInsets.top)
        }
        //清除之前所有布局属性
        attrsArr.removeAll()
        
        //开始创建每一个cell对应的布局属性
        let count:Int = self.collectionView?.numberOfItems(inSection: 0) ?? 0
        
        for i in 0 ..< count {
            let indexPath = IndexPath.init(item: i, section: 0)
            let attrs:UICollectionViewLayoutAttributes = self.layoutAttributesForItem(at: indexPath) ?? UICollectionViewLayoutAttributes.init()
            self.attrsArr.append(attrs)
        }
    }
    
    //返回对应index 位置cell的布局属性
    override func layoutAttributesForItem(at indexPath: IndexPath) -> UICollectionViewLayoutAttributes? {
        //创建布局属性
        let attrs = UICollectionViewLayoutAttributes.init(forCellWith: indexPath)
        //collectionView的宽度
        let collectionViewW:CGFloat = self.collectionView?.frame.size.width ?? 0


        // 设置布局属性的frame
        
        let cellW:CGFloat = (collectionViewW - self.edgeInsets.left - self.edgeInsets.right - (self.colunmCount-1).cgFloat * self.columnMargin) / self.colunmCount.cgFloat

        let cellH = self.delegate?.heightForItemAtIndexPath(indexPath: indexPath.item, itemWidth: cellW, waterFallLayout: self)
        
        //找最短的那一列
        var destColumn = 0
        var minColumnHeight = self.columnHeights[0]
        for i in 0..<LFDefaultColunmCount {
            let columnHeight = self.columnHeights[i]
            if minColumnHeight > columnHeight {
                minColumnHeight = columnHeight
                destColumn = i
            }
        }
        let cellX = self.edgeInsets.left + destColumn.cgFloat * (cellW + self.columnMargin)
        var cellY = minColumnHeight
        
        if cellY != self.edgeInsets.top {
            cellY += self.rowMargin
        }
        attrs.frame = CGRect.init(x: cellX, y: cellY, width: cellW, height: cellH ?? 0)
        // 更新最短那一列的高度
        self.columnHeights[destColumn] = attrs.frame.origin.y + attrs.frame.size.height
        // 记录内容的高度 - 即最长那一列的高度
        let  maxColumnHeight = self.columnHeights[destColumn]
        if self.contentHeight < maxColumnHeight {
            self.contentHeight = maxColumnHeight
        }
        
        return attrs
    }
    
//    * 决定cell的高度
    override func layoutAttributesForElements(in rect: CGRect) -> [UICollectionViewLayoutAttributes]? {
        return self.attrsArr as? [UICollectionViewLayoutAttributes]
    }
    
//    * 内容的高度
     
    override var collectionViewContentSize: CGSize {
        get {
            return CGSize.init(width: 0, height: self.contentHeight+self.edgeInsets.bottom)
        }
    }


}