How to add cache for solution to LeetCode Sum of Distances in Tree?

I am working on 834. Sum of Distances in Tree:

There is an undirected connected tree with n nodes labeled from 0 to n - 1 and n - 1 edges.

You are given the integer n and the array edges where edges[i] = [ai, bi] indicates that there is an edge between nodes ai and bi in the tree.

Return an array answer of length n where answer[i] is the sum of the distances between the ith node in the tree and all other nodes.

I tried to come up with my own solution, but got a Time Limit Exceeded (TLE) error. I realised that I need to add a cache for the solution.

Here is my code:

/**
 * @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;
};

The normal cache pattern is we cached the result at the end of func. On top, if we see cache value match, then we return cache value. (code above)

I am not sure what should be the cache index here.

I understand that start point to a particular point can go forward and backward, hence cache index can be start_pt + ‘‘ + end_point (forward) AND end_point + ‘‘ + start_pt (backward), to cache a particular path, but I am not able to fit to the cache pattern in the code above.

One test case that fails (as it takes too much time to finish) has 30000 nodes (numbered 0 to 29999) where node 0 has an edge to every other node, making all other nodes leaves of the tree.

1

There are faster solutions, but to answer your question, you should first make sure that the recursive call returns information that does not depend on the path you have already followed. In particular, caching something that has dist in it, reflects information that comes from outside the scope of the recursive call. That is not very practical for memoizing.

So first change the logic so that you don’t pass dist, but return information that only depends on the edge you just traversed (from parent to node). Return:

  • the number of paths
  • the sum of the path lengths

Also, you don’t really use the origNode parameter, so we can drop that one as well:

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 = (parent, node) => {

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

        const arr = graph[node];

        const result = { numPaths: 0, sumPathLengths: 0 };
        for(let i=0; i<arr?.length; ++i) {
            const childnode = arr[i];

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

            const tmp = recur(node, childnode);
            result.numPaths += tmp.numPaths;
            result.sumPathLengths += tmp.sumPathLengths;
        }
        if (parent != node) { // If we have follewed an edge:
            result.numPaths += 1; // The path that stops at node
            result.sumPathLengths += result.numPaths; // ...add edge to all paths
        }

        //cache[?] = sum;

        return result;
    }

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

    return res;
};

Now we can think of memoization:

As the call depends on parent and node, the key for your cache is this pair. You could multiply one by n and then sum them up to get a unique key:

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 = (parent, node) => {
        const key = parent * n + node;

        if(cache[key]) {
            return cache[key];
        }

        const arr = graph[node];

        const result = { numPaths: 0, sumPathLengths: 0 };
        for(let i=0; i<arr?.length; ++i) {
            const childnode = arr[i];

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

            const tmp = recur(node, childnode);
            result.numPaths += tmp.numPaths;
            result.sumPathLengths += tmp.sumPathLengths;
        }
        if (parent != node) { // If we have follewed an edge:
            result.numPaths += 1; // The path that stops at node
            result.sumPathLengths += result.numPaths; // ...add edge to all paths
        }

        cache[key] = result;
        return result;
    }

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

    return res;
};

This code will still take too long because even if you have cached the values for the edges, this still means that all neighbors of a node need to be visited, which can be a lot for a tree with a high branching factor.

You could improve on this to also check the cached value you might already have for a node, i.e. with key node * node + node:

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 = (parent, node) => {
        const key = parent * n + node;

        if(cache[key]) {
            return cache[key];
        }
        const keyNode = node * n + node; // The key for the node (not the edge)
        if (cache[keyNode]) { // We have the info for the node and the opposite edge:
            const nodeResult = cache[keyNode];
            const keyOppositeEdge = node * n + parent;
            const edgeResult = cache[keyOppositeEdge];
            // Derive the info for the forward edge
            const numPaths = nodeResult.numPaths - edgeResult.numPaths + 1;
            const result = {
                numPaths, 
                sumPathLengths: nodeResult.sumPathLengths - edgeResult.sumPathLengths + numPaths
            };
            cache[key] = result;
            return result;
        }

        const arr = graph[node];

        const result = { numPaths: 0, sumPathLengths: 0 };
        for(let i=0; i<arr?.length; ++i) {
            const childnode = arr[i];

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

            const tmp = recur(node, childnode);
            result.numPaths += tmp.numPaths;
            result.sumPathLengths += tmp.sumPathLengths;
        }
        if (parent != node) { // If we have followed an edge:
            result.numPaths += 1; // The path that stops at node
            result.sumPathLengths += result.numPaths; // ...add edge to all paths
        }

        cache[key] = result;
        return result;
    }

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

    return res;
};

This runs in an acceptable time, but there is a lot of caching involved, which has its overhead.

Another approach

The more common approach is to perform two DFS traversals from a chosen root (like node 0), one to accumulate the “downward” counts/sums, and one to accumulate the “upward” counts/sums, using the results from the first traversal.

For instance like this:

var sumOfDistancesInTree = function (n, edges) {
    const graph = Array.from({length: n}, () => []);
    const numPaths = Array(n).fill(1);
    const sumPathLengths = Array(n).fill(0);

    for (const [node1, node2] of edges) {
        graph[node1].push(node2);
        graph[node2].push(node1);
    }

    function dfs1(parent, node) {
        for (const child of graph[node]) {
            if (child === parent) {
                continue;
            }
            dfs1(node, child);
            numPaths[node] += numPaths[child];
            sumPathLengths[node] += sumPathLengths[child] + numPaths[child];
        }
    }

    function dfs2(parent, node) {
        for (const child of graph[node]) {
            if (child === parent) {
                continue;
            }
            sumPathLengths[child] = sumPathLengths[node] - 2 * numPaths[child] + n;
            dfs2(node, child);
        }
    }

    dfs1(0, 0);
    dfs2(0, 0);
    return sumPathLengths;
};

Trang chủ Giới thiệu Sinh nhật bé trai Sinh nhật bé gái Tổ chức sự kiện Biểu diễn giải trí Dịch vụ khác Trang trí tiệc cưới Tổ chức khai trương Tư vấn dịch vụ Thư viện ảnh Tin tức - sự kiện Liên hệ Chú hề sinh nhật Trang trí YEAR END PARTY công ty Trang trí tất niên cuối năm Trang trí tất niên xu hướng mới nhất Trang trí sinh nhật bé trai Hải Đăng Trang trí sinh nhật bé Khánh Vân Trang trí sinh nhật Bích Ngân Trang trí sinh nhật bé Thanh Trang Thuê ông già Noel phát quà Biểu diễn xiếc khỉ Xiếc quay đĩa Dịch vụ tổ chức sự kiện 5 sao Thông tin về chúng tôi Dịch vụ sinh nhật bé trai Dịch vụ sinh nhật bé gái Sự kiện trọn gói Các tiết mục giải trí Dịch vụ bổ trợ Tiệc cưới sang trọng Dịch vụ khai trương Tư vấn tổ chức sự kiện Hình ảnh sự kiện Cập nhật tin tức Liên hệ ngay Thuê chú hề chuyên nghiệp Tiệc tất niên cho công ty Trang trí tiệc cuối năm Tiệc tất niên độc đáo Sinh nhật bé Hải Đăng Sinh nhật đáng yêu bé Khánh Vân Sinh nhật sang trọng Bích Ngân Tiệc sinh nhật bé Thanh Trang Dịch vụ ông già Noel Xiếc thú vui nhộn Biểu diễn xiếc quay đĩa Dịch vụ tổ chức tiệc uy tín Khám phá dịch vụ của chúng tôi Tiệc sinh nhật cho bé trai Trang trí tiệc cho bé gái Gói sự kiện chuyên nghiệp Chương trình giải trí hấp dẫn Dịch vụ hỗ trợ sự kiện Trang trí tiệc cưới đẹp Khởi đầu thành công với khai trương Chuyên gia tư vấn sự kiện Xem ảnh các sự kiện đẹp Tin mới về sự kiện Kết nối với đội ngũ chuyên gia Chú hề vui nhộn cho tiệc sinh nhật Ý tưởng tiệc cuối năm Tất niên độc đáo Trang trí tiệc hiện đại Tổ chức sinh nhật cho Hải Đăng Sinh nhật độc quyền Khánh Vân Phong cách tiệc Bích Ngân Trang trí tiệc bé Thanh Trang Thuê dịch vụ ông già Noel chuyên nghiệp Xem xiếc khỉ đặc sắc Xiếc quay đĩa thú vị
Trang chủ Giới thiệu Sinh nhật bé trai Sinh nhật bé gái Tổ chức sự kiện Biểu diễn giải trí Dịch vụ khác Trang trí tiệc cưới Tổ chức khai trương Tư vấn dịch vụ Thư viện ảnh Tin tức - sự kiện Liên hệ Chú hề sinh nhật Trang trí YEAR END PARTY công ty Trang trí tất niên cuối năm Trang trí tất niên xu hướng mới nhất Trang trí sinh nhật bé trai Hải Đăng Trang trí sinh nhật bé Khánh Vân Trang trí sinh nhật Bích Ngân Trang trí sinh nhật bé Thanh Trang Thuê ông già Noel phát quà Biểu diễn xiếc khỉ Xiếc quay đĩa
Thiết kế website Thiết kế website Thiết kế website Cách kháng tài khoản quảng cáo Mua bán Fanpage Facebook Dịch vụ SEO Tổ chức sinh nhật