昨天我的代码,有一个 Bug,就是我没有处理好嵌套的数组元素为空的情况,我写了一个简单的 TestCase,大家也可以试试自己的代码是否处理好了这种情况:
- (void)testEmptyArray {
NSArray *arr = @[ @[ @[ ]], @[@[ @[ @[ ]]]]];
NSArrayIterator *c = [[NSArrayIterator alloc] initWithArray:arr];
XCTAssertEqualObjects(nil, [c next]);
XCTAssertEqualObjects(nil, [c next]);
}
于是乎,我发现我的代码可以再优化一下,用递归的方式来处理空数组的逻辑似乎是写起来更简单的,于是我优化之后的逻辑如下:
- 判断栈是否为空,如果为空则返回 nil。
- 从栈中取出元素,看是否遍历到了结尾,如果是的话,则出栈。
- 判断第 2 步是否使栈为空,如果为空,则返回 nil。
- 终于拿到元素了,这一步判断拿到的元素是否是数组。
- 如果是数组,则重新生成一个遍历的 NSArrayIteratorCursor 对象,放到栈中,并且递归调用自己。
- 如果不是数组,就把元素返回,同时更新索引到下一个位置。
整个代码也变得更短更清楚了一些,如下所示:
next 方法的实现:
- (id)next {
// 1. 判断栈是否为空,如果为空则返回 nil。
if (_stack.count == 0) {
return nil;
}
// 2. 从栈中取出元素,看是否遍历到了结尾,如果是的话,则出栈。
NSArrayIteratorCursor *c;
c = [_stack lastObject];
while (c.index == c.array.count && _stack.count > 0) {
[_stack removeLastObject];
c = [_stack lastObject];
}
// 3. 判断第2步是否使栈为空,如果为空,则返回 nil。
if (_stack.count == 0) {
return nil;
}
// 4. 终于拿到元素了,这一步判断拿到的元素是否是数组。
id item = c.array[c.index];
if ([item isKindOfClass:[NSArray class]]) {
c.index++;
// 5. 如果是数组,则重新生成一个遍历的
// NSArrayIteratorCursor 对象,放到栈中, 然后递归调用 next 方法
[self setupStackWithArray:item];
return [self next];
}
// 6. 如果到了这一步,说明拿到了一个非数组的元素,这样就可以把元素返回,
// 同时更新索引到下一个位置。
c.index++;
return item;
}
初使化部分:
- (id)initWithArray:(NSArray *)array {
self = [super init];
if (self) {
_originArray = array;
_stack = [NSMutableArray array];
[self setupStackWithArray:array];
}
return self;
}
- (void)setupStackWithArray:(NSArray *)array {
NSArrayIteratorCursor *c = [[NSArrayIteratorCursor alloc] initWithArray:array];
[_stack addObject:c];
}
感谢读者 @狒狒 的指正!代码和 TestCase 也已经更新到 gist 上了:https://gist.github.com/tangqiaoboy/452e106e0472b9e90cf17de180b6d211