我是照著 MJ 的視訊敲得.謝謝 MJ老師
其中,每個方法的註釋,有一個數字,表示每一個方法的呼叫順序

#import <UIKit/UIKit.h>
@class WaterFlowLayout;
@protocol WaterFlowLayoutDelegate <NSObject>

- (CGFloat)waterflowLayout:(WaterFlowLayout *)waterflowLayout heightForWidth:(CGFloat)width atIndexPath:(NSIndexPath *)indexPath;


@end

@interface
 
 WaterFlowLayout : UICollectionViewLayout
@property (nonatomic , assign) UIEdgeInsets sectionInset;
/** 每一列的間距 */
@property (nonatomic , assign) CGFloat columnMargin;
/** 每一行的間距 */
@property (nonatomic , assign) CGFloat rowMargin;
/** 顯示多少列 */
@property (nonatomic , assign) int columnsCount;

@property (nonatomic , weak) id<WaterFlowLayoutDelegate> delegate;
@end
 

#import "WaterFlowLayout.h"

@interface WaterFlowLayout ()
/** 這個字典用來儲存每一列最大的 Y值 */
@property (nonatomic ,  strong) NSMutableDictionary *maxYDict;
/** 存放所有的佈局屬性 */
@property (nonatomic , strong) NSMutableArray *attrsArray;
@end

@implementation WaterFlowLayout

- (NSMutableDictionary *)maxYDict
{
    if
 
 (!_maxYDict) {
        _maxYDict = [[NSMutableDictionary alloc] init];

    }
    return _maxYDict;
}

- (NSMutableArray *)attrsArray
{
    if (!_attrsArray) {
        _attrsArray = [NSMutableArray array];
    }
    return _attrsArray;
}

- (instancetype)init
{
    self = [super init];
    if (self) {
        self.rowMargin = 10;
        self.columnMargin = 10;
        self.sectionInset = UIEdgeInsetsMake(10, 10, 10, 10);
        /** 預設顯示 */
        self.columnsCount = 3;
    }
    return self;
}

/** 每一次滑動的時候呼叫 */
- (BOOL)shouldInvalidateLayoutForBoundsChange:(CGRect)newBounds
{
    return YES;
}

// 1.每次佈局之前的準備
- (void)prepareLayout
{
    [super prepareLayout];
    // 1.清空最大的 Y 值
    for (int i = 0; i < self.columnsCount; i++) {
        NSString *column = [NSString stringWithFormat:@"%d" , i];
        self.maxYDict[column] = @(self.sectionInset.top);
    }
    // 2.假設所有 cell 的屬性
    [self.attrsArray removeAllObjects];
    NSInteger count = [self.collectionView numberOfItemsInSection:0];
    for (int i = 0; i < count; i++) {
        UICollectionViewLayoutAttributes *attrs = [self layoutAttributesForItemAtIndexPath:[NSIndexPath indexPathForItem:i inSection:0]];
        [self.attrsArray addObject:attrs];
    }
}

/**
 *  3.返回所有的尺寸
 *  會多次呼叫
 */
- (CGSize)collectionViewContentSize
{
    __block NSString *maxColumn = @"0";
    [self.maxYDict enumerateKeysAndObjectsUsingBlock:^(NSString *column, NSNumber *maxY, BOOL *stop) {
        if ([maxY floatValue] > [self.maxYDict[maxColumn] floatValue]) {
            maxColumn = column;
        }
    }];
    return CGSizeMake(0, [self.maxYDict[maxColumn] floatValue] + self.sectionInset.bottom);
}

/**
 *  2.返回 indexPath 這個位置 item 的佈局屬性
 *  會多次呼叫
 */
- (UICollectionViewLayoutAttributes *)layoutAttributesForItemAtIndexPath:(NSIndexPath *)indexPath
{
    // 假設最短的那一列是第0列
    __block NSString *minColumn = @"0";
    [self.maxYDict enumerateKeysAndObjectsUsingBlock:^(NSString *column, NSNumber *maxY, BOOL *stop) {
        if ([maxY floatValue] < [self.maxYDict[minColumn] floatValue]) {
            minColumn = column;
        }
    }];

    // 計算寬度
    CGFloat width = (self.collectionView.frame.size.width - self.sectionInset.left - self.sectionInset.right - (self.columnsCount - 1) * self.columnMargin) / self.columnsCount;

    CGFloat height = [self.delegate waterflowLayout:self heightForWidth:width atIndexPath:indexPath];

    // 計算位置
    CGFloat x = self.sectionInset.left + (width + self.columnMargin) * [minColumn intValue];
    CGFloat y = [self.maxYDict[minColumn] floatValue] + self.rowMargin;

    // 更新這一列的最大 Y值
    self.maxYDict[minColumn] = @(y + height);
    // 建立屬性
    UICollectionViewLayoutAttributes *attrs = [UICollectionViewLayoutAttributes layoutAttributesForCellWithIndexPath:indexPath];
    attrs.frame = CGRectMake(x, y, width, height);
    return attrs;
}

// 4.返回 rect 範圍內的佈局屬性
- (NSArray<UICollectionViewLayoutAttributes *> *)layoutAttributesForElementsInRect:(CGRect)rect
{
    return self.attrsArray;
}
@end