diff --git a/NHBalancedFlowLayout/NHBalancedFlowLayout.m b/NHBalancedFlowLayout/NHBalancedFlowLayout.m
index aa08265..db44fa6 100644
--- a/NHBalancedFlowLayout/NHBalancedFlowLayout.m
+++ b/NHBalancedFlowLayout/NHBalancedFlowLayout.m
@@ -104,10 +104,10 @@ - (void)prepareLayout
     [self clearItemFrames];
     
     // create new item frame sections
-    _numberOfItemFrameSections = [self.collectionView numberOfSections];
+    _numberOfItemFrameSections = [self.collectionView.dataSource numberOfSectionsInCollectionView:self.collectionView];
     _itemFrameSections = (CGRect **)malloc(sizeof(CGRect *) * _numberOfItemFrameSections);
-    
-    for (int section = 0; section < [self.collectionView numberOfSections]; section++) {
+
+    for (int section = 0; section < _numberOfItemFrameSections; section++) {
         // add new item frames array to sections array
         NSInteger numberOfItemsInSections = [self.collectionView numberOfItemsInSection:section];
         CGRect *itemFrames = (CGRect *)malloc(sizeof(CGRect) * numberOfItemsInSections);
@@ -167,14 +167,16 @@ - (CGSize)collectionViewContentSize
 - (NSArray *)layoutAttributesForElementsInRect:(CGRect)rect
 {
     NSMutableArray *layoutAttributes = [NSMutableArray array];
-    
-    for (NSInteger section = 0, n = [self.collectionView numberOfSections]; section < n; section++) {
+    NSInteger n = [self.collectionView.dataSource numberOfSectionsInCollectionView:self.collectionView];
+
+    for (NSInteger section = 0; section < n; section++) {
         NSIndexPath *sectionIndexPath = [NSIndexPath indexPathForItem:0 inSection:section];
 
         UICollectionViewLayoutAttributes *headerAttributes = [self layoutAttributesForSupplementaryViewOfKind:UICollectionElementKindSectionHeader
                                                                                                   atIndexPath:sectionIndexPath];
         
-        if (! CGSizeEqualToSize(headerAttributes.frame.size, CGSizeZero) && CGRectIntersectsRect(headerAttributes.frame, rect)) {
+        CGSize size = headerAttributes.frame.size;
+        if (size.height != 0 && size.width != 0 && CGRectIntersectsRect(headerAttributes.frame, rect)) {
             [layoutAttributes addObject:headerAttributes];
         }
             
@@ -188,8 +190,8 @@ - (NSArray *)layoutAttributesForElementsInRect:(CGRect)rect
         
         UICollectionViewLayoutAttributes *footerAttributes = [self layoutAttributesForSupplementaryViewOfKind:UICollectionElementKindSectionFooter
                                                                                                   atIndexPath:sectionIndexPath];
-        
-        if (! CGSizeEqualToSize(footerAttributes.frame.size, CGSizeZero) && CGRectIntersectsRect(footerAttributes.frame, rect)) {
+        size = footerAttributes.frame.size;
+        if (size.width != 0 && size.height != 0 && CGRectIntersectsRect(footerAttributes.frame, rect)) {
             [layoutAttributes addObject:footerAttributes];
         }
     }
@@ -253,7 +255,9 @@ - (CGRect)footerFrameForSection:(NSInteger)section
 - (CGFloat)totalItemSizeForSection:(NSInteger)section preferredRowSize:(CGFloat)preferredRowSize
 {
     CGFloat totalItemSize = 0;
-    for (NSInteger i = 0, n = [self.collectionView numberOfItemsInSection:section]; i < n; i++) {
+    NSUInteger n = [self.collectionView.dataSource collectionView:self.collectionView numberOfItemsInSection:section];
+    
+    for (NSInteger i = 0; i < n; i++) {
         CGSize preferredSize = [self.delegate collectionView:self.collectionView layout:self preferredSizeForItemAtIndexPath:[NSIndexPath indexPathForItem:i inSection:section]];
         
         if (self.scrollDirection == UICollectionViewScrollDirectionVertical) {
@@ -282,6 +286,12 @@ - (NSArray *)weightsForItemsInSection:(NSInteger)section
 - (void)setFrames:(CGRect *)frames forItemsInSection:(NSInteger)section numberOfRows:(NSUInteger)numberOfRows sectionOffset:(CGPoint)sectionOffset sectionSize:(CGSize *)sectionSize
 {
     NSArray *weights = [self weightsForItemsInSection:section];
+    
+    if (weights.count == 0) {
+        *sectionSize = CGSizeZero;
+        return;
+    }
+
     NSArray *partition = [NHLinearPartition linearPartitionForSequence:weights numberOfPartitions:numberOfRows];
     
     int i = 0;