iOS学习（面试题）：唐巧 - 实现一个嵌套数组的迭代器

给你一个嵌套的 NSArray 数据，实现一个迭代器类，该类提供一个 next() 方法，可以依次的取出这个 NSArray 中的数据。

比如 NSArray 如果是 [1,[4,3],6,[5,[1,0]]]， 则最终应该输出：1, 4, 3, 6, 5, 1, 0 。

另外，实现一个 allObjects 方法，可以一次性取出所有元素。

先说第二问吧，第二问比较简单：实现一个 allObjects 方法，可以一次性取出所有元素。

对于此问，我们可以实现一个递归函数，在函数中判断数组中的元素是否又是数组，如果是的话，就递归调用自己，如果不是数组，则加入到一个 NSMutableArray 中即可。

下面是示例代码：

- (NSArray *)allObjects {
    NSMutableArray *result = [NSMutableArray array];
    [self fillArray:_originArray into:result];
    return result;
}
- (void)fillArray:(NSArray *)array into:(NSMutableArray *)result {
    for (NSInteger i = 0; i < array.count; i++) {
        if ([array[i] isKindOfClass:[NSArray class]]) {
            [self fillArray:array[i] into:result];
        } else {
            [result addObject:array[i]];
        }
    }
}

如果你还在纠结掌握递归有什么意义的话，欢迎翻翻我半年前写的另一篇文章：递归的故事（上），递归的故事（下）。

接下来让我们来看第一问，在同学的回复中，我看到很多人用第二问的办法，把数组整个另外保存一份，然后再记录一个下标，每次返回其中一个。这个方法当然是可行的，但是大部分的迭代器通常都不会这么实现。因为这么实现的话，数组需要整个复制一遍，空间复杂度是 O(N)。

所以，我个人认为本题第一问更好的解法是：
记录下遍历的位置，然后每次遍历时更新位置。由于本题中元素是一个嵌套数组，所以我们为了记录下位置，就需要两个变量：一个是当前正在遍历的子数组，另一个是这个数组遍历到的位置。

我在实现的时候，定义了一个名为 NSArrayIteratorCursor 的类来记录这些内容，NSArrayIteratorCursor 的定义和实现如下：

@interface NSArrayInteratorCursor : NSObject

@property (nonatomic, strong) NSArray *array;
@property (nonatomic, assign) NSUInteger index;

@end

@implementation NSArrayInteratorCursor

- (id)initWithArray:(NSArray *)array {
    if (self = [super init]) {
        self.array = array;
        self.index = 0;
    }
    return self;
}

@end

由于数组在遍历的时候可能产生递归，就像我们实现 allObjects 方法那样。所以我们需要处理递归时的 NSArrayIteratorCursor 的保存，我在实现的时候，拿数组当作栈，来实现保存遍历时的状态。

最终，我实现了一个迭代器类，名字叫 NSArrayIterator，用于最终提供 next 方法的实现。这个类有两个私有变量，一个是刚刚说的那个栈，另一个是原数组的引用。

@interface NSArrayIterator : NSObject

- (id)initWithArray:(NSArray *)array;
- (id)next;
- (NSArray *)allObjects;

@end

@implementation NSArrayIterator {
    NSMutableArray *_stack;
    NSArray *_originArray;
}

- (id)initWithArray:(NSArray *)array {
    if (self = [super init]) {
        _originArray = array;
        _stack = [NSMutableArray array];
    }
    return self;
}

- (void)setupStack {
    NSArrayInteratorCursor *c = [[NSArrayInteratorCursor alloc] initWithArray:_originArray];
    [_stack addObject:c];
}

@end

接下来就是最关键的代码了，即实现 next 方法，在 next 方法的实现逻辑中，我们需要：

1. 判断栈是否为空，如果为空则返回 nil。
2. 从栈中取出元素，看是否遍历到了结尾，如果是的话，则出栈。
3. 判断第 2 步是否使栈为空，如果为空，则返回 nil。
4. 终于拿到元素了，这一步判断拿到的元素是否是数组。
5. 如果是数组，则重新生成一个遍历的 NSArrayIteratorCursor 对象，放到栈中。
6. 重新从栈中拿出第一个元素，循环回到第 4 步的判断。
7. 如果到了这一步，说明拿到了一个非数组的元素，这样就可以把元素返回，同时更新索引到下一个位置。

以下是相关的代码，对于没有算法基础的同学，可能读起来还是比较累，其实我写起来也不快，所以希望你能多理解一下，其实核心思想就是手工操作栈的入栈和出栈：

 

- (id)next {
    // 1.判断栈是否为空，如果为空则返回 nil。
    if ([_stack count] == 0) {
        return nil;
    }
    // 2.从栈中取出元素，看是否遍历到了结尾，如果是的话，则出栈。
    NSArrayIteratorCursor *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];
    while ([item isKindOfClass:[NSArray class]]) {
        c.index++;
        // 5. 如果是数组，则重新生成一个遍历的 NSArrayIteratorCursor 对象，放到栈中。
        NSArrayIteratorCursor *nc = [[NSArrayIteratorCursor alloc] initWithArray:item];
        [_stack addObject:nc];
        // 6. 重新从栈中拿出第一个元素，循环回到第 4 步的判断。
        c = nc;
        item = c.array[c.index];
    }
    // 7. 如果到了这一步，说明拿到了一个非数组的元素，这样就可以把元素返回，同时更新索引到下一个位置。
    c.index++;
    return item;
}

最终，我想说这个只是我个人想出来的解法，很可能不是最优的，甚至可能也有很多问题，比如，这个代码有很多可以进一步 challenge 的地方：

• 这个代码是线程安全的吗？如果我们要实现一个线程安全的迭代器，应该怎么做？
• 如果在使用迭代器的时候，数组被修改了，会怎么样？
• 如何检测在遍历元素的时候，数组被修改了？
• 如何避免在遍历元素的时候，数组被修改？

 

有一个 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. 终于拿到元素了，这一步判断拿到的元素是否是数组。
5. 如果是数组，则重新生成一个遍历的 NSArrayIteratorCursor 对象，放到栈中，并且递归调用自己。
6. 如果不是数组，就把元素返回，同时更新索引到下一个位置。

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

- (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;
}