XMGWaterflowLayout.h

#import <UIKit/UIKit.h>

@interface XMGWaterflowLayout : UICollectionViewLayout

@end

XMGWaterflowLayout.m

#import "XMGWaterflowLayout.h"

/** 默認(rèn)的列數(shù) */
static const NSInteger XMGDefaultColumnCount = 3;
/** 每一列之間的間距 */
static const CGFloat XMGDefaultColumnMargin = 10;
/** 每一行之間的間距 */
static const CGFloat XMGDefaultRowMargin = 10;
/** 邊緣間距 */
static const UIEdgeInsets XMGDefaultEdgeInsets = {10, 10, 10, 10};

@interface XMGWaterflowLayout()
/** 存放所有cell的布局屬性 */
@property (nonatomic, strong) NSMutableArray *attrsArray;
/** 存放所有列的當(dāng)前高度 */
@property (nonatomic, strong) NSMutableArray *columnHeights;
@end

@implementation XMGWaterflowLayout

- (NSMutableArray *)columnHeights
{
    if (!_columnHeights) {
        _columnHeights = [NSMutableArray array];
    }
    return _columnHeights;
}

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

/**
 * 初始化
 */
- (void)prepareLayout
{
    [super prepareLayout];
    
    // 清除以前計算的所有高度
    [self.columnHeights removeAllObjects];
    for (NSInteger i = 0; i < XMGDefaultColumnCount; i++) {
        [self.columnHeights addObject:@(XMGDefaultEdgeInsets.top)];
    }

    // 清除之前所有的布局屬性
    [self.attrsArray removeAllObjects];
    // 開始創(chuàng)建每一個cell對應(yīng)的布局屬性
    NSInteger count = [self.collectionView numberOfItemsInSection:0];
    for (NSInteger i = 0; i < count; i++) {
        // 創(chuàng)建位置
        NSIndexPath *indexPath = [NSIndexPath indexPathForItem:i inSection:0];
        // 獲取indexPath位置cell對應(yīng)的布局屬性
        UICollectionViewLayoutAttributes *attrs = [self layoutAttributesForItemAtIndexPath:indexPath];
        [self.attrsArray addObject:attrs];
    }
}

/**
 * 決定cell的排布
 */
- (NSArray *)layoutAttributesForElementsInRect:(CGRect)rect
{
    return self.attrsArray;
}

/**
 * 返回indexPath位置cell對應(yīng)的布局屬性
 */
- (UICollectionViewLayoutAttributes *)layoutAttributesForItemAtIndexPath:(NSIndexPath *)indexPath
{
    // 創(chuàng)建布局屬性
    UICollectionViewLayoutAttributes *attrs = [UICollectionViewLayoutAttributes layoutAttributesForCellWithIndexPath:indexPath];
    
    // collectionView的寬度
    CGFloat collectionViewW = self.collectionView.frame.size.width;
    
    // 設(shè)置布局屬性的frame
    CGFloat w = (collectionViewW - XMGDefaultEdgeInsets.left - XMGDefaultEdgeInsets.right - (XMGDefaultColumnCount - 1) * XMGDefaultColumnMargin) / XMGDefaultColumnCount;
    CGFloat h = 50 + arc4random_uniform(100);
    
    // 找出高度最短的那一列
//    __block NSInteger destColumn = 0;
//    __block CGFloat minColumnHeight = MAXFLOAT;
//    [self.columnHeights enumerateObjectsUsingBlock:^(NSNumber *columnHeightNumber, NSUInteger idx, BOOL *stop) {
//        CGFloat columnHeight = columnHeightNumber.doubleValue;
//        if (minColumnHeight > columnHeight) {
//            minColumnHeight = columnHeight;
//            destColumn = idx;
//        }
    //    }];
    
    // 找出高度最短的那一列
    NSInteger destColumn = 0;
    CGFloat minColumnHeight = [self.columnHeights[0] doubleValue];
    for (NSInteger i = 1; i < XMGDefaultColumnCount; i++) {
        // 取得第i列的高度
        CGFloat columnHeight = [self.columnHeights[i] doubleValue];
        
        if (minColumnHeight > columnHeight) {
            minColumnHeight = columnHeight;
            destColumn = i;
        }
    }
    
    CGFloat x = XMGDefaultEdgeInsets.left + destColumn * (w + XMGDefaultColumnMargin);
    CGFloat y = minColumnHeight;
    if (y != XMGDefaultEdgeInsets.top) {
        y += XMGDefaultRowMargin;
    }
    attrs.frame = CGRectMake(x, y, w, h);
    
    // 更新最短那列的高度
    self.columnHeights[destColumn] = @(CGRectGetMaxY(attrs.frame));
    
    return attrs;
}

- (CGSize)collectionViewContentSize
{
    CGFloat maxColumnHeight = [self.columnHeights[0] doubleValue];
    for (NSInteger i = 1; i < XMGDefaultColumnCount; i++) {
        // 取得第i列的高度
        CGFloat columnHeight = [self.columnHeights[i] doubleValue];
        
        if (maxColumnHeight < columnHeight) {
            maxColumnHeight = columnHeight;
        }
    }
    return CGSizeMake(0, maxColumnHeight + XMGDefaultEdgeInsets.bottom);
}

@end

ViewController.m

#import "ViewController.h"
#import "XMGWaterflowLayout.h"

@interface ViewController () <UICollectionViewDataSource>

@end

@implementation ViewController

static NSString * const XMGShopId = @"shop";

- (void)viewDidLoad {
    [super viewDidLoad];
    
    // 創(chuàng)建布局
    XMGWaterflowLayout *layout = [[XMGWaterflowLayout alloc] init];
    
    // 創(chuàng)建CollectionView
    UICollectionView *collectionView = [[UICollectionView alloc] initWithFrame:self.view.bounds collectionViewLayout:layout];
    collectionView.dataSource = self;
    [self.view addSubview:collectionView];
    
    // 注冊
    [collectionView registerClass:[UICollectionViewCell class] forCellWithReuseIdentifier:XMGShopId];
}
#pragma mark - <UICollectionViewDataSource>
- (NSInteger)collectionView:(UICollectionView *)collectionView numberOfItemsInSection:(NSInteger)section
{
    return 50;
}

- (UICollectionViewCell *)collectionView:(UICollectionView *)collectionView cellForItemAtIndexPath:(NSIndexPath *)indexPath
{
    UICollectionViewCell *cell = [collectionView dequeueReusableCellWithReuseIdentifier:XMGShopId forIndexPath:indexPath];
    
    cell.backgroundColor = [UIColor orangeColor];

    NSInteger tag = 10;
    UILabel *label = (UILabel *)[cell.contentView viewWithTag:tag];
    if (label == nil) {
        label = [[UILabel alloc] init];
        label.tag = tag;
        [cell.contentView addSubview:label];
    }

    label.text = [NSString stringWithFormat:@"%zd", indexPath.item];
    [label sizeToFit];
    
    return cell;
}
@end