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

问题

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

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

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

解答

本题的代码稍长，完整的代码我放在 gist 上了：https://gist.github.com/tangqiaoboy/452e106e0472b9e90cf17de180b6d211 ，以下是讲解。

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

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

- (NSArray *)allObjects {
    NSMutableArray *result = [NSMutableArray array];
    [self fillArray:_originArray into:result];
    return result;
}

- (void)fillArray:(NSArray *)array into:(NSMutableArray *)result {
    for (NSUInteger 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 NSArrayIteratorCursor : NSObject

@property (nonatomic) NSArray *array;
@property (nonatomic) NSUInteger index;

@end

@implementation NSArrayIteratorCursor

- (id)initWithArray:(NSArray *)array {
    self = [super init];
    if (self) {
        _array = array;
        _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 {
    self = [super init];
    if (self) {
        _originArray = array;
        _stack = [NSMutableArray array];
        [self setupStack];
    }
    return self;
}

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

接下来就是最关键的代码了，即实现 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;
    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;
}

在读者回复中，听榆大叔 和 yiplee 同学用了类似的做法，他们的代码在：

• https://gist.github.com/NathanLi/bb8e5240ffc3eab1f97b38726549efc2
• https://gist.github.com/yiplee/a20f0b9ca6e09dcdfe26b129afb7bec8

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

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

如果大家有想出更好的解法，欢迎留言告诉我。

PS：本系列的第一篇文章《寻找最近公共 View》被一个读者提供了更优的解法，时间复杂度 O(N)，空间复杂度只有 O(1)，解法非常有意思，我会另外撰文分享。

---

下一期的面试题：我们知道 block 默认是不能被取消掉的，请你封装一个可以被取消执行的 block wrapper 类，它的定义如下：

typedef void (^Block)();

@interface CancelableObject : NSObject

- (id)initWithBlock:(Block)block;
- (void)start;
- (void)cancel;

@end