LeetCode 3241 - Time Taken to Mark All Nodes

LeetCode Problem 3241

Difficulty: 🔴 Hard
Topics: Dynamic Programming, Tree, Depth-First Search, Graph Theory

Solution

Problem Understanding

We are given an undirected tree with n nodes. A tree is a connected graph with exactly n - 1 edges and no cycles. Each node has a special propagation delay determined entirely by its parity:

• Odd-numbered nodes become marked 1 time unit after one of their neighbors is marked.
• Even-numbered nodes become marked 2 time units after one of their neighbors is marked.

For every possible starting node i, we begin by marking only node i at time 0. The marking process then spreads through the tree according to the delay rule above.

The important detail is that the delay depends on the node being marked, not the node that spreads the mark.

Suppose we move from node u to node v.

• If v is odd, then v becomes marked at time[u] + 1.
• If v is even, then v becomes marked at time[u] + 2.

For each starting node, we must determine the earliest time at which every node in the tree has become marked. The answer for node i is therefore the maximum marking time among all nodes when propagation starts from i.

The input consists of:

• edges, a list of undirected edges
• The number of nodes is inferred as len(edges) + 1

The output is an array times where:

times[i] = total time needed to mark the entire tree when starting from node i

The constraints are large:

2 <= n <= 100000

This immediately rules out any solution that performs a full traversal from every node in quadratic time.

Because the graph is a tree, there is exactly one simple path between any two nodes. This property is extremely important because it means the marking time from one node to another is uniquely determined by the path between them.

Several edge cases are important:

• A tree with only two nodes.
• A star-shaped tree where one center connects to many leaves.
• A long chain, where propagation delays accumulate heavily.
• Mixed parity paths, where odd and even delays alternate.
• Starting from an even node versus starting from an odd node.

The problem guarantees the graph is always a valid tree, so we never need to handle disconnected components or cycles.

Approaches

Brute Force Approach

The most direct solution is to independently simulate the propagation process from every starting node.

For each node i:

1. Run a shortest-path traversal from i.
2. Moving into node v costs:

• 1 if v is odd
• 2 if v is even

3. Compute the shortest marking time to every node.
4. The answer for i is the maximum distance discovered.

Because edge weights are positive and either 1 or 2, Dijkstra's algorithm works correctly.

This approach is correct because the marking process behaves exactly like shortest-path propagation on a weighted tree.

However, running Dijkstra from every node is far too expensive.

• One traversal costs O(n log n)
• Repeating it for all nodes costs O(n^2 log n)

With n = 100000, this is infeasible.

Key Insight

The graph is a tree, which means every pair of nodes has exactly one path between them.

The marking time between two nodes is simply the sum of node-entry costs along that unique path.

This transforms the problem into:

For every node:
find the maximum weighted distance to any other node.

This is a classic rerooting DP problem on trees.

We can solve it using two DFS passes:

1. First DFS computes the best downward distance inside each subtree.
2. Second DFS reroots the tree and computes the best distance coming from outside the subtree.

For every node:

answer[node] =
max(best downward path, best upward/outside path)

This reduces the complexity to linear time.

Approach Comparison

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(n² log n)	O(n)	Run Dijkstra from every node
Optimal	O(n)	O(n)	Two-pass rerooting DP on tree

Algorithm Walkthrough

Step 1: Build the Adjacency List

Because the graph is a tree, an adjacency list is the most efficient representation.

For every edge [u, v]:

• Add v to u
• Add u to v

This gives us O(n) storage and traversal time.

Step 2: Define Transition Cost

When moving from node u to neighbor v, the propagation delay depends on v.

Define:

cost(v) = 1 if v is odd else 2

This represents the additional time needed for v to become marked.

Step 3: First DFS, Compute Downward Distances

We define:

down[u] =
maximum time needed to reach any node inside u's subtree,
starting from u

For every child v:

candidate = cost(v) + down[v]

We take the maximum over all children.

This DFS computes the best path going downward from each node.

While doing this, we also store:

• the best child contribution
• the second-best child contribution

These values are necessary for rerooting later.

Step 4: Second DFS, Compute Upward Distances

Now we compute:

up[u] =
best distance reaching nodes outside u's subtree

When moving from parent u to child v, we must determine the best path available that does not reuse v's subtree.

There are two possibilities:

1. A path coming from above u
2. A path through another child subtree of u

We choose the larger one.

If v contributed the maximum downward path of u, then we must use the second-best value instead.

Then:

up[v] =
cost(u) + best_available

because moving into u incurs cost(u).

Step 5: Compute Final Answers

For every node:

answer[u] = max(down[u], up[u])

This represents the farthest weighted distance from u to any node in the tree.

That is exactly the total time needed to mark the entire tree.

Why it works

The tree contains exactly one path between every pair of nodes. Therefore, the marking time between two nodes is uniquely determined by the sum of entry costs along that path.

The first DFS computes the longest reachable distance inside each subtree. The second DFS propagates information about paths outside the subtree. Together, they guarantee that every possible path starting from a node is considered exactly once.

Thus, for every node, we correctly compute the maximum marking time to any other node.

Python Solution

from typing import List

class Solution:
    def timeTaken(self, edges: List[List[int]]) -> List[int]:
        n = len(edges) + 1

        graph = [[] for _ in range(n)]

        for u, v in edges:
            graph[u].append(v)
            graph[v].append(u)

        def cost(node: int) -> int:
            return 1 if node % 2 == 1 else 2

        down = [0] * n

        best1 = [0] * n
        best2 = [0] * n

        def dfs1(node: int, parent: int) -> None:
            for nei in graph[node]:
                if nei == parent:
                    continue

                dfs1(nei, node)

                value = cost(nei) + down[nei]

                if value > best1[node]:
                    best2[node] = best1[node]
                    best1[node] = value
                elif value > best2[node]:
                    best2[node] = value

            down[node] = best1[node]

        dfs1(0, -1)

        up = [0] * n

        def dfs2(node: int, parent: int) -> None:
            for nei in graph[node]:
                if nei == parent:
                    continue

                use = best1[node]

                candidate = cost(nei) + down[nei]

                if candidate == best1[node]:
                    use = best2[node]

                best_outside = max(up[node], use)

                up[nei] = cost(node) + best_outside

                dfs2(nei, node)

        dfs2(0, -1)

        answer = [0] * n

        for i in range(n):
            answer[i] = max(down[i], up[i])

        return answer

The implementation begins by constructing the adjacency list representation of the tree.

The cost() helper function encapsulates the propagation rule. Moving into an odd node costs 1, while moving into an even node costs 2.

The first DFS computes the longest downward distance for every node. While processing children, we track both the largest and second-largest child contributions. This is essential because rerooting must sometimes exclude the current child's own subtree contribution.

The second DFS propagates information from parent to child. For each child, we determine the best path available outside that child's subtree. This may come either from another sibling subtree or from ancestors above the current node.

Finally, the answer for each node is the larger of:

• the best downward distance
• the best upward distance

This represents the farthest node reachable from that starting node.

Go Solution

package main

func timeTaken(edges [][]int) []int {
	n := len(edges) + 1

	graph := make([][]int, n)

	for _, e := range edges {
		u, v := e[0], e[1]
		graph[u] = append(graph[u], v)
		graph[v] = append(graph[v], u)
	}

	cost := func(node int) int {
		if node%2 == 1 {
			return 1
		}
		return 2
	}

	down := make([]int, n)
	best1 := make([]int, n)
	best2 := make([]int, n)

	var dfs1 func(int, int)

	dfs1 = func(node int, parent int) {
		for _, nei := range graph[node] {
			if nei == parent {
				continue
			}

			dfs1(nei, node)

			value := cost(nei) + down[nei]

			if value > best1[node] {
				best2[node] = best1[node]
				best1[node] = value
			} else if value > best2[node] {
				best2[node] = value
			}
		}

		down[node] = best1[node]
	}

	dfs1(0, -1)

	up := make([]int, n)

	var dfs2 func(int, int)

	dfs2 = func(node int, parent int) {
		for _, nei := range graph[node] {
			if nei == parent {
				continue
			}

			use := best1[node]

			candidate := cost(nei) + down[nei]

			if candidate == best1[node] {
				use = best2[node]
			}

			bestOutside := up[node]
			if use > bestOutside {
				bestOutside = use
			}

			up[nei] = cost(node) + bestOutside

			dfs2(nei, node)
		}
	}

	dfs2(0, -1)

	answer := make([]int, n)

	for i := 0; i < n; i++ {
		if down[i] > up[i] {
			answer[i] = down[i]
		} else {
			answer[i] = up[i]
		}
	}

	return answer
}

The Go implementation follows the same rerooting DP structure as the Python version.

Slices are used for adjacency lists and DP arrays. Go does not support nested recursive functions as naturally as Python, so DFS functions are declared as function variables before assignment.

Integer overflow is not a concern because the maximum path length is at most 2 * (n - 1), which easily fits within Go's int.

Worked Examples

Example 1

edges = [[0,1],[0,2]]

Tree:

    1
    |
    0
    |
    2

Node costs:

Node	Cost
0	2
1	1
2	2

First DFS

For node 1:

down[1] = 0

For node 2:

down[2] = 0

For node 0:

From child 1:

1 + down[1] = 1

From child 2:

2 + down[2] = 2

So:

Node	down
0	2
1	0
2	0

Second DFS

For child 1:

• Best sibling contribution = 2
• up[1] = cost(0) + 2 = 4

For child 2:

• Best sibling contribution = 1
• up[2] = cost(0) + 1 = 3

Final answers:

Node	max(down, up)
0	2
1	4
2	3

Output:

[2,4,3]

Example 2

edges = [[0,1]]

First DFS

down[1] = 0
down[0] = 1

Second DFS

up[1] = cost(0) = 2

Final:

Node	Answer
0	1
1	2

Output:

[1,2]

Example 3

edges = [[2,4],[0,1],[2,3],[0,2]]

Tree:

    1
    |
    0
    |
    2
   / \
  4   3

Costs

Node	Cost
0	2
1	1
2	2
3	1
4	2

Downward DP

Node	down
1	0
3	0
4	0
2	2
0	4

Upward DP

Node	up
0	0
1	6
2	3
3	5
4	5

Final Answers

Node	Answer
0	4
1	6
2	3
3	5
4	5

Output:

[4,6,3,5,5]

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n)	Each edge is processed a constant number of times across two DFS traversals
Space	O(n)	Adjacency list and DP arrays each require linear storage

The algorithm performs exactly two DFS traversals over the tree. Since a tree contains n - 1 edges, the total amount of work is linear in the number of nodes.

The auxiliary arrays down, up, best1, and best2 each store one value per node, resulting in linear memory usage.

Test Cases

from typing import List

class Solution:
    def timeTaken(self, edges: List[List[int]]) -> List[int]:
        n = len(edges) + 1

        graph = [[] for _ in range(n)]

        for u, v in edges:
            graph[u].append(v)
            graph[v].append(u)

        def cost(node: int) -> int:
            return 1 if node % 2 else 2

        down = [0] * n
        best1 = [0] * n
        best2 = [0] * n

        def dfs1(node: int, parent: int):
            for nei in graph[node]:
                if nei == parent:
                    continue

                dfs1(nei, node)

                value = cost(nei) + down[nei]

                if value > best1[node]:
                    best2[node] = best1[node]
                    best1[node] = value
                elif value > best2[node]:
                    best2[node] = value

            down[node] = best1[node]

        dfs1(0, -1)

        up = [0] * n

        def dfs2(node: int, parent: int):
            for nei in graph[node]:
                if nei == parent:
                    continue

                use = best1[node]

                candidate = cost(nei) + down[nei]

                if candidate == best1[node]:
                    use = best2[node]

                up[nei] = cost(node) + max(up[node], use)

                dfs2(nei, node)

        dfs2(0, -1)

        return [max(down[i], up[i]) for i in range(n)]

sol = Solution()

assert sol.timeTaken([[0,1],[0,2]]) == [2,4,3]  # example 1
assert sol.timeTaken([[0,1]]) == [1,2]  # example 2
assert sol.timeTaken([[2,4],[0,1],[2,3],[0,2]]) == [4,6,3,5,5]  # example 3

assert sol.timeTaken([[0,1],[1,2]]) == [3,2,3]  # simple chain
assert sol.timeTaken([[0,1],[1,2],[2,3]]) == [4,3,3,5]  # longer chain
assert sol.timeTaken([[0,1],[0,2],[0,3]]) == [2,4,3,4]  # star topology
assert sol.timeTaken([[1,0],[1,2],[1,3],[3,4]]) == [5,3,5,4,6]  # mixed parity tree
assert sol.timeTaken([[0,1],[1,2],[2,3],[3,4]]) == [6,5,3,5,6]  # deep path

Test Summary

Test	Why
[[0,1],[0,2]]	Validates basic branching behavior
[[0,1]]	Smallest valid tree
[[2,4],[0,1],[2,3],[0,2]]	Mixed parity and branching
[[0,1],[1,2]]	Simple chain structure
[[0,1],[1,2],[2,3]]	Longer propagation path
[[0,1],[0,2],[0,3]]	Star topology
[[1,0],[1,2],[1,3],[3,4]]	Uneven subtree depths
[[0,1],[1,2],[2,3],[3,4]]	Deep linear tree stress test

Edge Cases

Two-Node Tree

The smallest valid input contains exactly two nodes connected by one edge. This case is important because many tree algorithms accidentally assume the existence of grandchildren or sibling subtrees.

The implementation handles this naturally because both DFS traversals work correctly even when a node has only one neighbor.

Long Linear Chain

A path-shaped tree can create the largest possible propagation times because delays accumulate across many nodes.

This case is dangerous for inefficient algorithms because naive all-pairs traversal becomes quadratic. The rerooting DP solution still processes every edge only twice, maintaining linear complexity.

Multiple Children with Equal Best Paths

When several child subtrees produce the same maximum contribution, rerooting logic can become subtle. Incorrect implementations may accidentally reuse the current child's own path when computing upward values.

The implementation avoids this by storing both the best and second-best child contributions. If the current child provided the maximum path, we switch to the second-best value during rerooting.