From 4c424ba6505f4bdfe06848ba05a1b2414071444f Mon Sep 17 00:00:00 2001
From: Sherlock Yao <yaosherlock@gmail.com>
Date: Sat, 1 Mar 2014 20:50:53 +0800
Subject: [PATCH] add a override method for NSArray->sample: to return an array
 of random elements of given number count

---
 Classes/NSArray+ObjectiveSugar.h                    |  9 +++++++++
 Classes/NSArray+ObjectiveSugar.m                    | 13 +++++++++++++
 .../ObjectiveSugarTests/NSArrayCategoriesTests.m    | 11 +++++++++++
 3 files changed, 33 insertions(+)

diff --git a/Classes/NSArray+ObjectiveSugar.h b/Classes/NSArray+ObjectiveSugar.h
index ea7b5a0..08c667a 100644
--- a/Classes/NSArray+ObjectiveSugar.h
+++ b/Classes/NSArray+ObjectiveSugar.h
@@ -31,6 +31,15 @@
  */
 - (id)sample;
 
+/**
+ Take random `numberOfElements` out of the array, or the maximum amount of
+ elements if it is less.
+ 
+ @param Number of random elements to take from array
+ @return An array of elements
+ */
+- (NSArray *)sample:(NSUInteger)numberOfElements;
+
 /// Alias for -sample
 - (id)anyObject;
 
diff --git a/Classes/NSArray+ObjectiveSugar.m b/Classes/NSArray+ObjectiveSugar.m
index b2c2503..9ef3278 100644
--- a/Classes/NSArray+ObjectiveSugar.m
+++ b/Classes/NSArray+ObjectiveSugar.m
@@ -21,6 +21,19 @@ - (id)sample {
     return self[index];
 }
 
+- (NSArray *)sample:(NSUInteger)numberOfElements {
+  if (numberOfElements > self.count) {
+    numberOfElements = self.count;
+  }
+  NSMutableArray *array = [NSMutableArray arrayWithArray:self];
+  for (NSUInteger i = 0; i < numberOfElements; i++) {
+    NSInteger size = array.count - i;
+    NSInteger r = (arc4random() % size) + i;
+    [array exchangeObjectAtIndex:i withObjectAtIndex:r];
+  }
+  return [array subarrayWithRange:NSMakeRange(0, numberOfElements)];
+}
+
 - (id)objectForKeyedSubscript:(id)key {
     if ([key isKindOfClass:[NSString class]])
         return [self subarrayWithRange:[self rangeFromString:key]];
diff --git a/Example/ObjectiveSugarTests/NSArrayCategoriesTests.m b/Example/ObjectiveSugarTests/NSArrayCategoriesTests.m
index f3628f8..e118767 100644
--- a/Example/ObjectiveSugarTests/NSArrayCategoriesTests.m
+++ b/Example/ObjectiveSugarTests/NSArrayCategoriesTests.m
@@ -29,6 +29,17 @@
         NSArray *emptyArray = @[];
         [emptyArray.sample shouldBeNil];
     });
+  
+    it(@"-sample: returns an array of random elements with given count", ^{
+      NSArray *result = [sampleArray sample:2];
+      [[result should] haveCountOf:2];
+      [[sampleArray should] containObjectsInArray:result];
+    });
+  
+    it(@"-sample: returns an array of maximum amount when it is less than given count" , ^{
+      [[[sampleArray sample:12] should] haveCountOf:3];
+      [[[sampleArray sample:-1] should] haveCountOf:3];
+    });
 
     context(@"Iterating using block", ^{