如何为LeetCode 834题「树中距离之和」解法添加缓存？

题目描述 #

给定一棵包含n个节点（编号0到n-1）和n-1条边的无向连通树。给定整数n和边数组edges，其中edges[i] = [a_i, b_i]表示树中节点a_i和b_i之间存在一条边。返回长度为n的数组answer，其中answer[i]是第i个节点与树中所有其他节点的距离之和。

遇到的问题 #

我自己实现的解法出现了Time Limit Exceeded（TLE）错误，尝试添加缓存优化时，不确定该用什么作为缓存键。常规的start_pt + '_' + end_point这类缓存键无法适配当前代码的递归逻辑。其中一个超时测试用例包含30000个节点，节点0与其他所有节点相连，其余节点均为叶子节点。

当前实现代码 #

/**
 * @param {number} n
 * @param {number[][]} edges
 * @return {number[]}
 */
var sumOfDistancesInTree = function(n, edges) {
    // * g: graph
    const graph = {};
    let res = [];
    let cache = {};

    // * g: build graph
    for(let i=0; i<edges.length; ++i) {
        const ind1 = parseInt(edges[i][0]);
        const ind2 = parseInt(edges[i][1]);

        graph[ind1] = graph[ind1] || [];
        graph[ind1].push(ind2);

        graph[ind2] = graph[ind2] || [];
        graph[ind2].push(ind1);
    }

    const recur = (origNode, parent, node, dist) => {

        // if(cache[?]) {
        //     return cache[?]
        // }

        const arr = graph[node];

        let sum = dist;
        for(let i=0; i<arr?.length; ++i) {
            const childnode = arr[i];

            if(parent === childnode) {
                continue;
            }

            const tmp = recur(origNode, node, childnode, dist+1);
            sum = sum + tmp;
        }

        // * g: try to add cache, but not sure what should be the cache index
        //cache[?] = sum;

        return sum;
    }

    // * g: graph is done
    for(let i=0; i<n; ++i) {
        const out = recur(i, i, i, 0);
        res.push(out);
    }

    return res;
};

问题分析与优化方案 #

为什么单纯加缓存没用 #

当前代码时间复杂度是O(n²)，对每个节点都做了一次全树遍历。即使给recur函数加缓存，缓存键需要标识从origNode到node的累计距离和，也就是用${origNode}_${node}作为键，但这样缓存条目数会达到O(n²)，对于n=3e4的情况，内存会直接溢出，遍历时间也无法有效降低，依然会超时。

正确的优化思路：树形DP（两次DFS） #

要把时间复杂度降到O(n)，需要用两次深度优先搜索：

1. 第一次DFS（后序遍历）：计算每个节点作为根时，其子树内所有节点到它的距离和sumDist，以及子树的节点数count。
   • 对于节点u，sumDist[u] = sum(sumDist[v] + count[v])，其中v是u的子节点。因为子树v的每个节点到u的距离都比到v多1，所以要加上count[v]。
2. 第二次DFS（前序遍历）：利用父节点的结果推导子节点的结果。
   • 当从父节点u移动到子节点v时，v子树内的count[v]个节点到v的距离减少1，其余n-count[v]个节点到v的距离增加1。因此sumDist[v] = sumDist[u] - count[v] + (n - count[v])。

优化后的代码 #

/**
 * @param {number} n
 * @param {number[][]} edges
 * @return {number[]}
 */
var sumOfDistancesInTree = function(n, edges) {
    // 构建邻接表
    const graph = Array.from({ length: n }, () => []);
    for (const [a, b] of edges) {
        graph[a].push(b);
        graph[b].push(a);
    }

    const sumDist = new Array(n).fill(0); // 存储每个节点的距离和
    const count = new Array(n).fill(1);   // 存储每个节点的子树节点数（包括自身）

    // 第一次DFS：后序遍历，计算sumDist和count
    const postDFS = (u, parent) => {
        for (const v of graph[u]) {
            if (v === parent) continue;
            postDFS(v, u);
            count[u] += count[v];
            sumDist[u] += sumDist[v] + count[v];
        }
    };

    // 第二次DFS：前序遍历，推导其他节点的sumDist
    const preDFS = (u, parent) => {
        for (const v of graph[u]) {
            if (v === parent) continue;
            sumDist[v] = sumDist[u] - count[v] + (n - count[v]);
            preDFS(v, u);
        }
    };

    postDFS(0, -1);
    preDFS(0, -1);

    return sumDist;
};

关于缓存的补充 #

如果一定要在原递归思路上加缓存，缓存键可以用${origNode}_${node}，因为recur(origNode, parent, node, dist)的返回值是node子树内所有节点到origNode的距离之和（含当前node到origNode的dist）。但这种方法在n较大时依然会超时，不推荐使用。

内容的提问来源于stack exchange，提问作者kenpeter