Groups from NSArray

GroupTests.m

//
//  GroupTests.m
//
//  Created by Ray Fix on 3/15/14.
//

#import <XCTest/XCTest.h>
#import "NSArray+Groups.h"

@interface GroupTests : XCTestCase
@end

@implementation GroupTests

- (void)testGroupsOfOneFromPopulationOfThree
{
  NSArray* groups = [@[@"A", @"B", @"C"] groupsOfSize:1];
  NSArray* answer = @[@[@"A"], @[@"B"], @[@"C"]];
  XCTAssertEqualObjects(groups, answer);
}

- (void)testGroupsOfTwoFromPopulationOfThree
{
  NSArray* groups = [@[@"A", @"B", @"C"] groupsOfSize:2];
  NSArray* answer = @[@[@"A",@"B"],@[@"A",@"C"],@[@"B",@"C"]];
  XCTAssertEqualObjects(groups, answer);
}

- (void)testGroupsOfThreeFromPopulationOfThree
{
  NSArray* groups = [@[@"A", @"B", @"C"] groupsOfSize:3];
  NSArray* answer = @[@[@"A",@"B",@"C"]];
  XCTAssertEqualObjects(groups, answer);
}

- (void)testGroupsOfThreeFromFromPopulationOfFive
{
  NSArray* groups = [@[@"A", @"B", @"C", @"D", @"E"] groupsOfSize:3];
  NSArray* answer = @[@[@"A",@"B",@"C"],
                      @[@"A",@"B",@"D"],
                      @[@"A",@"B",@"E"],
                      @[@"A",@"C",@"D"],
                      @[@"A",@"C",@"E"],
                      @[@"A",@"D",@"E"],
                      @[@"B",@"C",@"D"],
                      @[@"B",@"C",@"E"],
                      @[@"B",@"D",@"E"],
                      @[@"C",@"D",@"E"]];
  XCTAssertEqualObjects(groups, answer);
}

@end

NSArray+Groups.m

//
//  NSArray+Groups.m
//
//  Created by Ray Fix on 3/15/14.
//

#import "NSArray+Groups.h"

@implementation NSArray (Groups)

- (NSArray *)groupsOfSize:(int)n
{
  NSAssert(n > 0 && n <= self.count,
           @"Illegal group size n=%d array.count = %ld",n,(long)self.count);
  
  NSMutableArray* answer = [NSMutableArray array];
  
  if (n == 1) {
    for (id member in self) {
      [answer addObject:@[member]];
    }
    return answer;
  }
  
  for (int index = 0 ; index <= self.count - n; ++index) {
    id firstObject = self[index];
    
    NSRange remainingPopulationRange = NSMakeRange(index+1, self.count-index-1);
    NSArray* remainingPopulation = [self subarrayWithRange:remainingPopulationRange];
    
    NSArray* groups = [remainingPopulation groupsOfSize:n-1];
    
    for (NSArray* group in groups) {
      NSMutableArray* combinedGroup = [@[firstObject] mutableCopy];
      [combinedGroup addObjectsFromArray:group];
      [answer addObject:combinedGroup];
    }
  }
  return answer;
}
@end