← 返回首页

数组拉平：从递归到生成器

数组拉平问题的三种解法：递归、手写迭代器、生成器（yield*）。

发布 2019年3月14日 标签 #javascript #generators #iterators

~/posts/array-flatten-recursion-to-generators $ cat post.md

/ 语言 EN / 中文
/ 主题 / /

数组拉平

Array.prototype.flat 现在已经可以用了。它不算什么底层新技术，只是规范层面的实现。所谓”数组拉平”的故事是这样的：有一个数组，元素是 T 或 T[]，T 可以继续嵌套——

type SomeArray<T> = Array<T | SomeArray<T>>;

上面这段类型会报错，TS 不允许循环嵌套自身。

要做的就是把整个结构铺平成 T[]。

递归

最直接的实现是递归：

const flat = <T>(arr: SomeArray<T>): T[] => {
    return arr.reduce((previous: T[], current: T | SomeArray<T>) => {
        if (!Array.isArray(current)) {
            return previous.concat(current);
        }
        return previous.concat(flat(current));
    }, [] as T[]);
};

这段代码用了 SomeArray 仅作说明，类型本身在 TS 里不合法，但逻辑能跑。

没有花哨结构：当前元素是数组就递归，不是数组就直接拼到结果上。这里用 reduce + concat 串成的形式正好可以引到下一步——迭代。

可迭代

TS / ES 中的数组类型实现了 Iterable<T>。这背后是 迭代协议：值如果实现了协议，就可以被 for...of 或 [...something] 解构。协议要求的是一个 next() 方法，返回 { done, value }。

这套机制天然适合数组拉平——只要让 next() 在遇到 T[] 时递归进入子元素就行。

文档里的例子大致是：

var myIterator = {
    next: function () {
        // ...
    },
    [Symbol.iterator]: function () {
        return this;
    },
};

https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Iteration_protocols

用迭代器实现一个简化版本的 flat：

const flat = (arr) => {
    const rest = [...arr]; // clone
    const iterator = {
        [Symbol.iterator]: () => ({
            next: () => {
                if (rest.length === 0) {
                    return { done: true };
                }
                const current = rest.shift();
                if (Array.isArray(current)) {
                    rest.unshift(...current);
                } else {
                    rest.unshift(current);
                }
                return { value: rest.shift(), done: false };
            },
        }),
    };
    return [...iterator];
};

写出来比递归丑。TS 还更不让你这么写，强制要 class extends Iterator<T>，进一步劝退。好消息是还有更好的方案——生成器。

生成器

生成器的语法是 function*，见 Iterators and Generators。用 yield 关键字可以写一个无限自增的例子：

function* infiniteNumber(): IterableIterator<number> {
    let index: number = 0;
    while (1) {
        yield index++;
    }
}
const iterator: IterableIterator<number> = infiniteNumber();

可以理解为：函数运行到 yield 时暂停；下一次 next() 被调用时从暂停处继续；函数体跑完时返回 done: true。

把递归交给另一个生成器

yield* 关键字会把当前生成器接下来若干次 next() 转交给另一个生成器：

function* gen() {
    yield* ["a", "b", "c"];
}

https://developer.mozilla.org/en-US/docs/Web/JavaScript/Guide/Iterators_and_Generators

也就是说在 yield* 处，接下来的三次迭代由 ["a", "b", "c"] 的迭代器负责。利用这个机制，生成器里可以写出递归般的效果。

生成器版的 flat

写法和递归没什么本质区别，重点是用的时候记得解构。解构生成器时如果不小心容易跑出无限循环，最好先克隆一份数组、再在闭包里访问。

function* flat<T>(arr: Array<T | T[]>): IterableIterator<T> {
    const rest: Array<T | T[]> = [...arr]; // clone
    while (rest.length > 0) {
        const current: T | T[] = rest.shift();
        if (Array.isArray(current)) {
            yield* flat(current);
        } else {
            yield current;
        }
    }
}

console.log([...flat([1, [2, [4, 5]], 3])]);

到这一步，生成器和 yield* 的组合就把这件事写得相对清爽了。

← 返回首页