Skip to content

2277. Closest Node to Path in Tree 👍

  • Time: $O(n^2 + mn)$
  • Space: $O(n^2)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
class Solution {
 public:
  vector<int> closestNode(int n, vector<vector<int>>& edges,
                          vector<vector<int>>& query) {
    vector<int> ans;
    vector<vector<int>> tree(n);
    vector<vector<int>> dist(n, vector<int>(n, -1));

    for (const vector<int>& edge : edges) {
      const int u = edge[0];
      const int v = edge[1];
      tree[u].push_back(v);
      tree[v].push_back(u);
    }

    for (int i = 0; i < n; ++i)
      fillDist(tree, i, i, 0, dist);

    for (const vector<int>& query : queries) {
      const int start = query[0];
      const int end = query[1];
      const int node = query[2];
      ans.push_back(findClosest(tree, dist, start, end, node, start));
    }

    return ans;
  }

 private:
  void fillDist(const vector<vector<int>>& tree, int start, int u, int d,
                vector<vector<int>>& dist) {
    dist[start][u] = d;
    for (const int v : tree[u])
      if (dist[start][v] == -1)
        fillDist(tree, start, v, d + 1, dist);
  }

  int findClosest(const vector<vector<int>>& tree,
                  const vector<vector<int>>& dist, int u, int end, int node,
                  int ans) {
    for (const int v : tree[u])
      if (dist[v][end] < dist[u][end])
        return findClosest(tree, dist, v, end, node,
                           dist[ans][node] < dist[v][node] ? ans : v);
    return ans;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
class Solution {
  public int[] closestNode(int n, int[][] edges, int[][] query) {
    int[] ans = new int[query.length];
    List<Integer>[] tree = new List[n];
    int[][] dist = new int[n][n];
    Arrays.stream(dist).forEach(A -> Arrays.fill(A, -1));

    for (int i = 0; i < n; ++i)
      tree[i] = new ArrayList<>();

    for (int[] edge : edges) {
      final int u = edge[0];
      final int v = edge[1];
      tree[u].add(v);
      tree[v].add(u);
    }

    for (int i = 0; i < n; ++i)
      fillDist(tree, i, i, 0, dist);

    for (int i = 0; i < query.length; ++i) {
      final int start = query[i][0];
      final int end = query[i][1];
      final int node = query[i][2];
      ans[i] = findClosest(tree, dist, start, end, node, start);
    }

    return ans;
  }

  private void fillDist(List<Integer>[] tree, int start, int u, int d, int[][] dist) {
    dist[start][u] = d;
    for (final int v : tree[u])
      if (dist[start][v] == -1)
        fillDist(tree, start, v, d + 1, dist);
  }

  private int findClosest(List<Integer>[] tree, int[][] dist, int u, int end, int node, int ans) {
    for (final int v : tree[u])
      if (dist[v][end] < dist[u][end])
        return findClosest(tree, dist, v, end, node, dist[ans][node] < dist[v][node] ? ans : v);
    return ans;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
class Solution:
  def closestNode(
      self,
      n: int,
      edges: list[list[int]],
      query: list[list[int]],
  ) -> list[int]:
    tree = [[] for _ in range(n)]
    dist = [[-1] * n for _ in range(n)]

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

    def fillDist(start: int, u: int, d: int) -> None:
      dist[start][u] = d
      for v in tree[u]:
        if dist[start][v] == -1:
          fillDist(start, v, d + 1)

    for i in range(n):
      fillDist(i, i, 0)

    def findClosest(u: int, end: int, node: int, ans: int) -> int:
      for v in tree[u]:
        if dist[v][end] < dist[u][end]:
          return findClosest(
              v, end, node, ans if dist[ans][node] < dist[v][node] else v)
      return ans

    return [findClosest(start, end, node, start)
            for start, end, node in query]