LeetCode 3313 - Find the Last Marked Nodes in Tree

We are given an undirected tree with n nodes labeled from 0 to n - 1. The tree is initially completely unmarked. We then simulate a spreading process that behaves like a multi-source breadth-first expansion: at each second, every unmarked node that has at least one already…

leetcodehardtreedepth-first-search

LeetCode Problem 3313

Difficulty: 🔴 Hard
Topics: Tree, Depth-First Search

Solution

Problem Understanding

We are given an undirected tree with n nodes labeled from 0 to n - 1. The tree is initially completely unmarked. We then simulate a spreading process that behaves like a multi-source breadth-first expansion: at each second, every unmarked node that has at least one already marked neighbor becomes marked.

The key twist is that we repeat this process for every possible starting node. For each node i, we assume that only node i is initially marked at time t = 0, then the marking spreads outward through edges one layer per second. Eventually, all nodes become marked. We must determine which node is the last to be marked for each starting node i. If multiple nodes are tied for being last, any one valid node can be returned.

The input is a tree, meaning it is connected and acyclic, so there is exactly one simple path between any pair of nodes. The constraint n <= 10^5 implies that any solution with quadratic behavior per node is impossible, and we must rely on structural properties of trees and reuse computations.

Important edge cases include very small trees such as n = 2, linear chains, and star-shaped trees. These extremes determine whether the solution correctly handles diameter endpoints and distance propagation. The problem guarantees connectivity and acyclicity, so we do not need to handle disconnected components or cycles.

Approaches

Brute Force Approach

For each starting node i, we simulate a BFS from i and compute the distance from i to every other node. The last marked node is simply the node with maximum distance from i. This is correct because the marking process expands exactly one edge per second, so the time at which a node is marked equals its shortest-path distance from the start node.

However, doing a BFS from every node leads to repeated traversal of the same tree structure, resulting in an O(n^2) time complexity, which is too slow for n = 10^5.

Key Insight for Optimal Solution

The crucial observation is that in a tree, the node farthest from a given start node is always one endpoint of a diameter path of the tree, or can be determined via distances to diameter endpoints.

A tree has a well-defined diameter, defined as the longest shortest path between any two nodes. Let the diameter endpoints be A and B. For any node i, the farthest node from i must lie at one of these two endpoints. This is a classical tree property: eccentricity in a tree is always achieved at a diameter endpoint.

Thus, instead of running BFS from every node, we:

  1. Find one endpoint of the diameter using BFS.
  2. Find the other endpoint using BFS from that node.
  3. Precompute distances from both endpoints to all nodes.
  4. For each node i, compare distA[i] and distB[i]. The node with larger distance determines which endpoint is farthest, and thus the last marked node.

This reduces the problem to constant-time per node after two BFS traversals.

Complexity Comparison

Approach Time Complexity Space Complexity Notes
Brute Force O(n^2) O(n) BFS from every node
Optimal O(n) O(n) Two BFS + distance comparison

Algorithm Walkthrough

  1. Build an adjacency list representation of the tree from the edge list. This allows efficient traversal of neighbors in O(1) amortized time per edge.
  2. Run a BFS from an arbitrary node, for example node 0, to find the farthest node from it. This node is one endpoint of the diameter, call it A. This works because BFS from any node in a tree reaches a diameter endpoint when selecting the farthest node.
  3. Run a second BFS starting from A to find the farthest node from A. This node is the other endpoint of the diameter, call it B.
  4. While performing the BFS from A, record distA[x], the shortest distance from A to every node x.
  5. Run another BFS from B to compute distB[x], the shortest distance from B to every node x.
  6. For each node i, determine which endpoint is farther: if distA[i] >= distB[i], then node A is at least as far as B, otherwise B is farther. The last marked node starting from i will be whichever endpoint is farther from i.
  7. Construct the answer array accordingly.

Why it works

In a tree, the eccentricity of a node (maximum distance to any other node) is always realized at one of the diameter endpoints. Since BFS from a node spreads uniformly along shortest paths, the last node to be marked is exactly the node farthest in shortest-path distance. Because all farthest nodes lie among diameter endpoints, comparing distances to the two endpoints is sufficient to determine the answer.

Python Solution

from typing import List
from collections import deque

class Solution:
    def lastMarkedNodes(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 bfs(start: int):
            dist = [-1] * n
            q = deque([start])
            dist[start] = 0
            farthest = start
            
            while q:
                node = q.popleft()
                for nei in graph[node]:
                    if dist[nei] == -1:
                        dist[nei] = dist[node] + 1
                        q.append(nei)
                        if dist[nei] > dist[farthest]:
                            farthest = nei
            return farthest, dist
        
        # Step 1: find one diameter endpoint
        A, _ = bfs(0)
        
        # Step 2: find other endpoint and distances from A
        B, distA = bfs(A)
        
        # Step 3: distances from B
        _, distB = bfs(B)
        
        # Step 4: decide answer for each node
        ans = [0] * n
        for i in range(n):
            ans[i] = A if distA[i] >= distB[i] else B
        
        return ans

Implementation Explanation

The solution begins by constructing an adjacency list for efficient traversal. The nested bfs function computes both the farthest node from a given start and the full distance array from that start.

We first use BFS from node 0 to find a candidate diameter endpoint A. Then we run BFS from A to obtain the true opposite endpoint B and simultaneously record distances from A. A third BFS from B computes distances from the other endpoint.

Finally, we compare distances from each node to A and B. Since the farthest reachable nodes in a tree are always among diameter endpoints, this comparison yields the correct last marked node.

Go Solution

package main

import "container/list"

func lastMarkedNodes(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)
	}

	bfs := func(start int) (int, []int) {
		dist := make([]int, n)
		for i := range dist {
			dist[i] = -1
		}

		q := list.New()
		q.PushBack(start)
		dist[start] = 0

		farthest := start

		for q.Len() > 0 {
			front := q.Remove(q.Front()).(int)
			for _, nei := range graph[front] {
				if dist[nei] == -1 {
					dist[nei] = dist[front] + 1
					q.PushBack(nei)
					if dist[nei] > dist[farthest] {
						farthest = nei
					}
				}
			}
		}

		return farthest, dist
	}

	A, _ := bfs(0)
	B, distA := bfs(A)
	_, distB := bfs(B)

	ans := make([]int, n)
	for i := 0; i < n; i++ {
		if distA[i] >= distB[i] {
			ans[i] = A
		} else {
			ans[i] = B
		}
	}

	return ans
}

Go-specific Notes

The Go implementation uses container/list for BFS queue operations since Go does not have a native deque. Distances are initialized explicitly to -1 to mark unvisited nodes. Slice allocation replaces Python lists, and careful type assertions are required when popping from the queue.

Worked Examples

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

Tree structure is a star centered at 0.

First BFS from 0 gives either 1 or 2 as endpoint A. Suppose A = 1. BFS from 1 gives 2 as B. Distances:

  • From A=1: distA = [1,0,2]
  • From B=2: distB = [1,2,0]

Comparison:

  • Node 0: tie, choose 1
  • Node 1: A is closer
  • Node 2: B is closer

Final answer: [2,2,1] depending on tie resolution.

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

Diameter endpoints are 1 and 3/4. BFS identifies endpoints A=1, B=3 (or 4 depending on traversal). Distances from both endpoints determine which is farther per node. Nodes near the right branch (3,4) prefer one endpoint, while left branch prefers the other.

Final result aligns with provided output [3,3,1,1,1].

Complexity Analysis

Measure Complexity Explanation
Time O(n) Three BFS traversals over a tree, each visiting nodes and edges once
Space O(n) Adjacency list and distance arrays

The linear complexity arises because each BFS is O(n) in a tree with n - 1 edges, and we perform a constant number of BFS passes.

Test Cases

# provided examples
assert Solution().lastMarkedNodes([[0,1],[0,2]]) in ([2,2,1], [1,1,2])  # tie allowed
assert Solution().lastMarkedNodes([[0,1]]) == [1,0]  # two node tree
assert Solution().lastMarkedNodes([[0,1],[0,2],[2,3],[2,4]])[0] in [3,4]  # structure correctness

# chain tree
assert Solution().lastMarkedNodes([[0,1],[1,2],[2,3]])[0] in [3]  # endpoint behavior

# star tree
assert len(Solution().lastMarkedNodes([[0,1],[0,2],[0,3],[0,4]])) == 5  # full coverage

# minimal edge case
assert Solution().lastMarkedNodes([[0,1]]) == [1,0]
Test Why
[[0,1],[0,2]] verifies star structure and tie handling
[[0,1]] smallest tree
chain tree verifies diameter endpoints
star tree ensures correctness in high branching
minimal edge case ensures base correctness

Edge Cases

One important edge case is a two-node tree. In this case, both nodes are diameter endpoints, and each node’s last marked node is trivially the other node. Incorrect handling of BFS initialization or tie-breaking can easily break this case if distances are not initialized properly.

Another edge case is a linear chain. Here every node lies on the diameter path, and endpoints are clearly the only valid farthest nodes. Any solution relying on incorrect assumptions about branching will still work here, but it is useful for validating correctness of diameter computation.

A third edge case is a star-shaped tree. The center node has equal distance to all leaves, and each leaf’s farthest node is another leaf. This case stresses tie-breaking logic when multiple nodes share maximum distance. The implementation handles this by allowing either endpoint when distances are equal, satisfying the problem’s “any valid answer” requirement.