【续】iOS 面试题（八）：实现一个嵌套数组的迭代器

昨天我的代码，有一个 Bug，就是我没有处理好嵌套的数组元素为空的情况，我写了一个简单的 TestCase，大家也可以试试自己的代码是否处理好了这种情况：

- (void)testEmptyArray {
    NSArray *arr = @[ @[ @[ ]], @[@[ @[ @[ ]]]]];
    NSArrayIterator *c = [[NSArrayIterator alloc] initWithArray:arr];
    XCTAssertEqualObjects(nil, [c next]);
    XCTAssertEqualObjects(nil, [c next]);
}

于是乎，我发现我的代码可以再优化一下，用递归的方式来处理空数组的逻辑似乎是写起来更简单的，于是我优化之后的逻辑如下：

1. 判断栈是否为空，如果为空则返回 nil。
2. 从栈中取出元素，看是否遍历到了结尾，如果是的话，则出栈。
3. 判断第 2 步是否使栈为空，如果为空，则返回 nil。
4. 终于拿到元素了，这一步判断拿到的元素是否是数组。
   1. 如果是数组，则重新生成一个遍历的 NSArrayIteratorCursor 对象，放到栈中，并且递归调用自己。
   2. 如果不是数组，就把元素返回，同时更新索引到下一个位置。

整个代码也变得更短更清楚了一些，如下所示：

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