The O(n) Club: Sum of Distances in Tree — Two DFSes, No Tears

The O(n) Club: Sum of Distances in Tree — Two DFSes, No Tears

⚡ TL;DR

Want to sum distances from every node to every other node? Don’t write a BFS-for-every-node death loop unless you’re benchmarking fans. Instead: two DFS passes, one for tallying up the minions, one for updating as you re-root. Short on drama, big on results:

// Java post-order DFS: gather subtree sums and sizes
private void postOrder(int node, int parent) {
    for (int child : tree[node]) if (child != parent) {
        postOrder(child, node);
        count[node] += count[child];
        res[node] += res[child] + count[child];
    }
}

🧠 How Devs Usually Mess This Up

New hire energy: “Let’s do a BFS/DFS from every node!” — lovely, unless n is >10, and you suddenly own a cryptocurrency heat generator. That’s O(n²) and terrible. The MVP failure mode here? People confuse subtree size (count[node]) with depth. Remember: count[node] is how big your cult following is, not how famous you are. Also, don’t recalculate answers for each root from scratch — tree DP is all about recycling.

🔍 What’s Actually Going On

This algorithm is what happens when dynamic programming and accounting have a baby. Imagine your company as a tree: you, as CEO, want to know your ‘distance’ to every employee. But you’re lazy. So you make each manager total up their subordinates’ scores and pass it up. Enter post-order DFS (bottom-up), where every node sums up its descendants.

Then, the magic: re-rooting! Pretend you make each child a CEO for the day. Using a simple trick, you give them their score in O(1), piggybacking off parent’s answer, using the formula:

ans[child] = ans[parent] - count[child] + (n - count[child]);

It’s corporate ladder climbing, without the performance reviews.

🛠️ PseudoCode

• Build the tree as an adjacency list in Java:

  for (int[] e : edges) {
      tree[e[0]].add(e[1]);
      tree[e[1]].add(e[0]);
  }

  If you don’t, you get an O(n²) LinkedList party. No one wants that.

• First DFS (post-order):

  postOrder(node, parent):
    count[node] = 1;
    for each child != parent:
      postOrder(child, node)
      count[node] += count[child]
      res[node] += res[child] + count[child]
  

  This sweeps up all the kids’ answers. Leaves start with count=1, res=0.

• Second DFS (pre-order):

  preOrder(node, parent):
    for each child != parent:
      res[child] = res[node] - count[child] + (n - count[child])
      preOrder(child, node)
  

  Suddenly, each child knows their sum, using their parent’s results like a true social climber.

• Return res[] — now res[i] is the answer for each node.

💻 The Code

class Solution {
    int[] res, count;
    List
<integer>[] tree;
    int n;
    public int[] sumOfDistancesInTree(int N, int[][] edges) {
        n = N;
        tree = new List[N];
        res = new int[N];
        count = new int[N];
        for (int i = 0; i < N; i++) tree[i] = new ArrayList<>();
        for (int[] e : edges) {
            tree[e[0]].add(e[1]);
            tree[e[1]].add(e[0]);
        }
        dfsPost(0, -1);
        dfsPre(0, -1);
        return res;
    }
    private void dfsPost(int node, int parent) {
        count[node] = 1;
        for (int child : tree[node]) {
            if (child == parent) continue;
            dfsPost(child, node);
            count[node] += count[child];
            res[node] += res[child] + count[child];
        }
    }
    private void dfsPre(int node, int parent) {
        for (int child : tree[node]) {
            if (child == parent) continue;
            res[child] = res[node] - count[child] + (n - count[child]);
            dfsPre(child, node);
        }
    }
}
</integer>

⚠️ Pitfalls, Traps & Runtime Smacks

• DFS Parent Check: If you forget it, enjoy that stack overflow.
• Subtree Size vs. Depth: Don’t mix these or you’ll have a philosophical crisis and a wrong answer.
• Edge Cases: Handle single-node “trees” and stars (one hub, all leaf workers). If you crash on these, the LeetCode bot will snark at you.
• Space: O(n) for arrays and lists; don’t pull out a HashMap unless you miss slow code.
• Time: The two DFSes are O(n). The only thing quadratic should be your caffeine intake.

🧠 Interviewer Brain-Tattoo

“Never BFS from every node in a tree unless you want an angry reviewer. Re-root, reuse, get O(n), go home early.”