具体实现如下：

LYFWaterflowLayout.h

#import <UIKit/UIKit.h>

@class LYFWaterflowLayout;

@protocol LYFWaterflowLayoutDelegate@required

/** cell高度  */

- (CGFloat)waterflowLayout:(LYFWaterflowLayout *)waterflowLayout heightForItemAtIndex:(NSUInteger)index itemWidth:(CGFloat)itemWidth;

@optional

/** 行数或列数（默认3  */

- (CGFloat)columnCountInWaterflowLayout:(LYFWaterflowLayout *)waterflowLayout;

/** 列间距 */

- (CGFloat)columnMarginInWaterflowLayout:(LYFWaterflowLayout *)waterflowLayout;

/** 行间距 */

- (CGFloat)rowMarginInWaterflowLayout:(LYFWaterflowLayout *)waterflowLayout;

/** 内边距 */

- (UIEdgeInsets)edgeInsetsInWaterflowLayout:(LYFWaterflowLayout *)waterflowLayout;

@end

@interface LYFWaterflowLayout : UICollectionViewLayout

/** 代理 */

@property (nonatomic, weak) iddelegate;

@end

LYFWaterflowLayout.m

/** 默认的列数 */

static const NSInteger LYFDefaultColumnCount = 3;

/** 每一列之间的间距 */

static const CGFloat LYFDefaultColumnMargin = 10;

/** 每一行之间的间距 */

static const CGFloat LYFDefaultRowMargin = 10;

/** 边缘间距 */

static const UIEdgeInsets LYFDefaultEdgeInsets = {10, 10, 10, 10};

@interface LYFWaterflowLayout()

/** 存放所有cell的布局属性 */

@property (nonatomic, strong) NSMutableArray *attrsArray;

/** 存放所有列的当前高度 */

@property (nonatomic, strong) NSMutableArray *columnHeights;

/** 内容的高度 */

@property (nonatomic, assign) CGFloat contentHeight;

- (CGFloat)rowMargin;

- (CGFloat)columnMargin;

- (NSInteger)columnCount;

- (UIEdgeInsets)edgeInsets;

@end

@implementation LYFWaterflowLayout

#pragma mark - 常见数据处理

- (CGFloat)rowMargin

{

if ([self.delegate respondsToSelector:@selector(rowMarginInWaterflowLayout:)]) {

return [self.delegate rowMarginInWaterflowLayout:self];

} else {

return LYFDefaultRowMargin;

}

}

- (CGFloat)columnMargin

{

if ([self.delegate respondsToSelector:@selector(columnMarginInWaterflowLayout:)]) {

return [self.delegate columnMarginInWaterflowLayout:self];

} else {

return LYFDefaultColumnMargin;

}

}

- (NSInteger)columnCount

{

if ([self.delegate respondsToSelector:@selector(columnCountInWaterflowLayout:)]) {

return [self.delegate columnCountInWaterflowLayout:self];

} else {

return LYFDefaultColumnCount;

}

}

- (UIEdgeInsets)edgeInsets

{

if ([self.delegate respondsToSelector:@selector(edgeInsetsInWaterflowLayout:)]) {

return [self.delegate edgeInsetsInWaterflowLayout:self];

} else {

return LYFDefaultEdgeInsets;

}

}

#pragma mark - 懒加载

- (NSMutableArray *)columnHeights

{

if (!_columnHeights) {

_columnHeights = [NSMutableArray array];

}

return _columnHeights;

}

- (NSMutableArray *)attrsArray

{

if (!_attrsArray) {

_attrsArray = [NSMutableArray array];

}

return _attrsArray;

}

/**

* 初始化

*/

- (void)prepareLayout

{

[super prepareLayout];

self.contentHeight = 0;

// 清除以前计算的所有高度

[self.columnHeights removeAllObjects];

for (NSInteger i = 0; i < self.columnCount; i++) {

[self.columnHeights addObject:@(self.edgeInsets.top)];

}

// 清除之前所有的布局属性

[self.attrsArray removeAllObjects];

// 开始创建每一个cell对应的布局属性

NSInteger count = [self.collectionView numberOfItemsInSection:0];

for (NSInteger i = 0; i < count; i++) {

// 创建位置

NSIndexPath *indexPath = [NSIndexPath indexPathForItem:i inSection:0];

// 获取indexPath位置cell对应的布局属性

UICollectionViewLayoutAttributes *attrs = [self layoutAttributesForItemAtIndexPath:indexPath];

[self.attrsArray addObject:attrs];

}

}

/**

* 决定cell的排布

*/

- (NSArray *)layoutAttributesForElementsInRect:(CGRect)rect

{

return self.attrsArray;

}

/**

* 返回indexPath位置cell对应的布局属性

*/

- (UICollectionViewLayoutAttributes *)layoutAttributesForItemAtIndexPath:(NSIndexPath *)indexPath

{

// 创建布局属性

UICollectionViewLayoutAttributes *attrs = [UICollectionViewLayoutAttributes layoutAttributesForCellWithIndexPath:indexPath];

// collectionView的宽度

CGFloat collectionViewW = self.collectionView.frame.size.width;

// 设置布局属性的frame

CGFloat w = (collectionViewW - self.edgeInsets.left - self.edgeInsets.right - (self.columnCount - 1) * self.columnMargin) / self.columnCount;

CGFloat h = [self.delegate waterflowLayout:self heightForItemAtIndex:indexPath.item itemWidth:w];

// 找出高度最短的那一列

NSInteger destColumn = 0;

CGFloat minColumnHeight = [self.columnHeights[0] doubleValue];

for (NSInteger i = 1; i < self.columnCount; i++) {

// 取得第i列的高度

CGFloat columnHeight = [self.columnHeights[i] doubleValue];

if (minColumnHeight > columnHeight) {

minColumnHeight = columnHeight;

destColumn = i;

}

}

CGFloat x = self.edgeInsets.left + destColumn * (w + self.columnMargin);

CGFloat y = minColumnHeight;

if (y != self.edgeInsets.top) {

y += self.rowMargin;

}

attrs.frame = CGRectMake(x, y, w, h);

// 更新最短那列的高度

self.columnHeights[destColumn] = @(CGRectGetMaxY(attrs.frame));

// 记录内容的高度

CGFloat columnHeight = [self.columnHeights[destColumn] doubleValue];

if (self.contentHeight < columnHeight) {

self.contentHeight = columnHeight;

}

return attrs;

}

- (CGSize)collectionViewContentSize

{

return CGSizeMake(0, self.contentHeight + self.edgeInsets.bottom);

}

具体使用如下：

LYFWaterflowLayout *layout = [[LYFWaterflowLayout alloc]init];

layout.delegate =self;

UICollectionView *collectionView = [[UICollectionView alloc]initWithFrame:self.view.bounds collectionViewLayout:layout]; 

cell高度设置(必须实现的代理)

- (CGFloat)waterflowLayout:(LYFWaterflowLayout *)waterflowLayout heightForItemAtIndex:(NSUInteger)index itemWidth:(CGFloat)itemWidth

{

return self.view.frame.size.height / 3;

}

设置行数或列数（默认为3）

- (CGFloat)columnCountInWaterflowLayout:(LYFWaterflowLayout *)waterflowLayout

{

return 4;

}

设置列间距（默认为10）

- (CGFloat)columnMarginInWaterflowLayout:(LYFWaterflowLayout *)waterflowLayout

{

return 20;

}

设置行间距（默认为10）

- (CGFloat)rowMarginInWaterflowLayout:(LYFWaterflowLayout *)waterflowLayout

{

return 20;

}

设置内边距（默认为10,10,10,10）

- (UIEdgeInsets)edgeInsetsInWaterflowLayout:(LYFWaterflowLayout *)waterflowLayout

{

return UIEdgeInsetsMake(20, 20, 20, 20);

}