使用 Typescript 实现 TreeMap

TreeMap 简介

首先看下列的 java 代码，你认为输入顺序会是什么？

import java.util.TreeMap;
import java.util.Set;;

public class UseTreeMap {
    public static void main(String[] args) {
        TreeMap<Integer, String> treeMap = new TreeMap<>();
        treeMap.put(5, "five");
        treeMap.put(3, "three");
        treeMap.put(7, "seven");
        Set<Integer> set = treeMap.keySet();
        for(Integer key: set) {
            System.out.println("treeMap value: " + treeMap.get(key));
        }
    }
}

以一个前端开发者看的话，应该是先 five ，然后 three ，最后 seven，可实际的顺序是先 three ，然后 five ，最后 seven 。原因是 java 的 TreeMap 会对传入的 key 进行大小排序。那么我们能不能使用 typescript 来实现一个 TreeMap 呢？答案当然是肯定的

实现增删改查

由于 Map 的使用类似于 class ，因此 TreeMap 也使用 class 来实现，首先我想到的是继承，至于调整 key 的顺序，我们在内部维护一个数组来记录 key ，并且由于 Map 的 key 是唯一的，数组的 key 也必须是唯一的。

class TreeMap<V> extends Map<number, V> {
  // 保证 key 的唯一性
  private readonly keysSet: number[] = [];
}

为什么不用 Set 这一数据结构呢？主要是数组提供了各种 api 方便对内部元素进行操作，而 Set 并没有

接着来考虑一个问题，哪些操作会对 Map 的 key 有影响，有三个 api : set 、delete 、clear 和 constructor ，其中 set 只有在 Map 中不存在传入的 key 时候才有影响

让我们先实现简单的 clear ，当调用该方便时，Map 的 size 也就为 0 了，那么相对应的，需要把 keysSet 清空

private clearKeys(): void {
    this.keysSet.length = 0
}

clear(): void {
    this.clearKeys();
    super.clear()
}

接着实现 set ，我们需要确保 keysSet 与 Map 的 set 在数量上和值是对应的，因此先判断 keysSet 有没有存在该 key ，存在就不需要操作了，然后根据情况插入到正确的位置

private insertKey(key: number): void {
    if (this.keysSet.includes(key)) {
        return
    }
    const index = this.keysSet.findIndex(k => k > key)
    if (index === -1) {
        this.keysSet.push(key)
    } else {
        this.keysSet.splice(index, 0, key)
    }
}

set(key: number, value: V): this {
    this.insertKey(key)
    return super.set(key, value)
}

再来实现 delete ，先判断要删除的 key 在 Map 中是否存在，存在执行删除操作

private deleteKey(key: number): void {
    const index = this.keysSet.findIndex(k => k === key)
    if (index !== -1) {
        this.keysSet.splice(index, 1)
    }
}

delete(key: number): boolean {
    if (this.has(key)) {
        this.deleteKey(key)
    }
    return super.delete(key)
}

最后再来实现 constructor ，其实可看做是初始进行了多次 set 操作，实现也很简单

constructor(entries?: readonly (readonly [number, V])[] | null) {
    super()
    for (const [key, value] of (entries ?? [])) {
        // 注意，这里不能使用 super.set ，而必须是 this.set ，因为需要对 keysSet 进行操作
        this.set(key, value)
    }
}

自此，我们已经实现了增删改的操作，还差个 get ，但 get 的逻辑与 Map 的逻辑并无任何不同，也不会对 keysSet 产生任何影响，因此可以不用实现

实现迭代功能

迭代首先想到的 forEach ，由于 keysSet 是有序且肯定有值的，因此可以直接遍历 keysSet，将值放入传入的回调函数就行了

forEach(callbackfn: (value: V, key: number, map: Map<number, V>) => void): void {
    const keysSet = this.keysSet
    for (let i = 0; i < keysSet.length; i++) {
        callbackfn(super.get(keysSet[i]) as V, keysSet[i], this)
    }
}

forEach 简单，但其他的就难了。以 keys 为例 ，这个方法需要返回 IterableIterator<number> ，那么能不能直接 return super.keys() 呢？

很遗憾，并不能，直接返回的话，key 的顺序是按插入的顺序来的，而不是正序返回，那么能不能手动实现接口呢？

这是 lib.es2015.iterable.d.ts 中的 IterableIterator 接口，实在是过于复杂，很显然，不好实现

interface SymbolConstructor {
    /**
     * A method that returns the default iterator for an object. Called by the semantics of the
     * for-of statement.
     */
    readonly iterator: unique symbol;
}

interface IteratorYieldResult<TYield> {
    done?: false;
    value: TYield;
}

interface IteratorReturnResult<TReturn> {
    done: true;
    value: TReturn;
}

type IteratorResult<T, TReturn = any> = IteratorYieldResult<T> | IteratorReturnResult<TReturn>;

interface Iterator<T, TReturn = any, TNext = undefined> {
    // NOTE: 'next' is defined using a tuple to ensure we report the correct assignability errors in all places.
    next(...args: [] | [TNext]): IteratorResult<T, TReturn>;
    return?(value?: TReturn): IteratorResult<T, TReturn>;
    throw?(e?: any): IteratorResult<T, TReturn>;
}

interface Iterable<T> {
    [Symbol.iterator](): Iterator<T>;
}

interface IterableIterator<T> extends Iterator<T> {
    [Symbol.iterator](): IterableIterator<T>;
}

但也不是没有曲线救国的方法，我在翻看上面那个 dts 文件的时候，也看到了下方这个

interface Array<T> {
    /**
     * Returns an iterable of values in the array
     */
    values(): IterableIterator<T>;
}

那么，TreeMap 的 keys 就好实现了，我们就有个现成的有序的 keysSet

keys(): IterableIterator<number> {
    return this.keysSet.values()
}

values 和 entries 的实现同理

values(): IterableIterator<V> {
    return this.keysSet.map(set => super.get(set) as V).values()
}

entries(): IterableIterator<[number, V]> {
    return this.keysSet.map((set: number): [number, V] => [set, super.get(set) as V]).values()
}

实现 Object.keys() 的效果

说实话，Map 的 keys 和 Object 的 keys 在使用方便程度上根本不是同一等级，既然我们用了 keysSet 作为顺序的标识，那么当然可以把它导出，只要有这种需求。需要注意的是，虽然我用了 readonly 来限定它，但是 push shift pop unshift 等方法仍然可以改变它，如果外部改变了 keysSet 的话，这意味着我们上述的迭代功能将全部失效。解决方法是在返回类型前面加一个 readonly 限定，这样外界将无法对 keysSet 做任何操作，包括我前面提到的那些数组方法。假如使用 javascript 来实现的话，就无法保证有效性了

getKeysSet(): readonly number[] {
    return this.keysSet
}

那么，Object.values() 和 Object.entries() 也是很容易实现的

单元测试

最后让我们来做个单测，由于 Map 中 key 的顺序是按插入顺序来的，我们按顺序升序插入 key ，它就应该符合 TreeMap 的特征，我们再不按升序顺序插入 key 在 TreeMap 中，然后比较两者的 api 的返回值，如果一致的话，说明我们的实现是成功的

// TreeMap.spec.ts
import TreeMap from "../data_structure/TreeMap";

const map = new Map<number, string>([[1, "first"]])
const treeMap = new TreeMap<string>([[5, "five"]])

map.set(3, "three").set(5, "five")
treeMap.set(3, "three").set(1, "first")

test("get", () => {
    [1, 3, 5, 7, 9].forEach(key => {
        expect(map.get(key)).toBe(treeMap.get(key))
    })
})

test("set", () => {
    map.set(7, "seven")
    treeMap.set(7, "seven")
    expect(map.get(7)).toBe(treeMap.get(7))
})

test("delete", () => {
    expect(treeMap.getKeysSet()).toEqual([1, 3, 5, 7])
    treeMap.delete(3)
    expect(treeMap.get(3)).toBe(undefined)
    expect(treeMap.getKeysSet()).toEqual([1, 5, 7])
    // 保持与 treeMap keys 数量一致
    map.delete(3)
})

test("forEach", () => {
    const kvs1: [number, string][] = []
    const kvs2: [number, string][] = []
    map.forEach((value, key) => {
        kvs1.push([key, value])
    })
    treeMap.forEach((value, key) => {
        kvs2.push([key, value])
    })
    expect(kvs1).toEqual(kvs2)
})

test("keys", () => {
    const iterator1 = map.keys()
    const iterator2 = treeMap.keys()
    expect(map.size).toBe(treeMap.size)
    for (let i = 0; i < map.size; i++) {
        expect(iterator1.next().value).toBe(iterator2.next().value)
    }
})

test("values", () => {
    const iterator1 = map.values()
    const iterator2 = treeMap.values()
    expect(map.size).toBe(treeMap.size)
    for (let i = 0; i < map.size; i++) {
        expect(iterator1.next().value).toBe(iterator2.next().value)
    }
})

test("entries", () => {
    const iterator1 = map.entries()
    const iterator2 = treeMap.entries()
    expect(map.size).toBe(treeMap.size)
    for (let i = 0; i < map.size; i++) {
        expect(iterator1.next().value).toEqual(iterator2.next().value)
    }
})

test("clear", () => {
    treeMap.clear()
    expect(treeMap.size).toBe(0)
    expect(treeMap.getKeysSet()).toEqual([])
})

下图是我跑的测试结果，全部通过

优化点

可以看到，上面的 TreeMap 实现只能是 number 作为 key 值，而且是升序的。如果你有兴趣的话，可以去实现让它能够降序，或者是其他的比如 string 作为 key 值，或者是任意的排序规则，我相信这不难。附上 TreeMap 代码 ，也欢迎在评论区互相探讨。