Skip to content

3486. Longest Special Path II

  • Time: $O(|V| + |E|)$
  • Space: $O(|V| + |E|)$
 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
48
49
50
51
52
53
54
55
class Solution {
 public:
  vector<int> longestSpecialPath(vector<vector<int>>& edges,
                                 vector<int>& nums) {
    vector<vector<pair<int, int>>> graph(nums.size());

    for (const vector<int>& edge : edges) {
      const int u = edge[0];
      const int v = edge[1];
      const int w = edge[2];
      graph[u].emplace_back(v, w);
      graph[v].emplace_back(u, w);
    }

    int maxLength = 0;
    int minNodes = 1;
    dfs(graph, 0, -1, /*leftBoundary=*/{0, 0}, /*prefix=*/{0},
        /*lastSeenDepth=*/{}, nums, maxLength, minNodes);
    return {maxLength, minNodes};
  }

 private:
  void dfs(const vector<vector<pair<int, int>>>& graph, int u, int prev,
           vector<int> leftBoundary, vector<int>&& prefix,
           unordered_map<int, int>&& lastSeenDepth, const vector<int>& nums,
           int& maxLength, int& minNodes) {
    const int prevDepth = lastSeenDepth[nums[u]];
    lastSeenDepth[nums[u]] = prefix.size();

    if (prevDepth != 0) {
      leftBoundary.push_back(prevDepth);
      ranges::sort(leftBoundary);
      leftBoundary = {leftBoundary[leftBoundary.size() - 2],
                      leftBoundary.back()};
    }

    const int length = prefix.back() - prefix[leftBoundary[0]];
    const int nodes = prefix.size() - leftBoundary[0];
    if (length > maxLength || (length == maxLength && nodes < minNodes)) {
      maxLength = length;
      minNodes = nodes;
    }

    for (const auto& [v, w] : graph[u]) {
      if (v == prev)
        continue;
      prefix.push_back(prefix.back() + w);
      dfs(graph, v, u, leftBoundary, std::move(prefix),
          std::move(lastSeenDepth), nums, maxLength, minNodes);
      prefix.pop_back();
    }

    lastSeenDepth[nums[u]] = prevDepth;
  }
};

 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
48
49
50
51
52
53
54
55
56
57
class Solution {
  public int[] longestSpecialPath(int[][] edges, int[] nums) {
    List<Pair<Integer, Integer>>[] graph = new List[nums.length];

    for (int i = 0; i < nums.length; ++i)
      graph[i] = new ArrayList<>();

    for (int[] edge : edges) {
      final int u = edge[0];
      final int v = edge[1];
      final int w = edge[2];
      graph[u].add(new Pair<>(v, w));
      graph[v].add(new Pair<>(u, w));
    }

    int[] ans = new int[2];
    ans[0] = 0; // maxLength
    ans[1] = 1; // minNodes
    dfs(graph, 0, -1, /*leftBoundary=*/List.of(0, 0),
        /*prefix=*/new ArrayList<>(List.of(0)),
        /*lastSeenDepth=*/new HashMap<>(), nums, ans);
    return ans;
  }

  private void dfs(List<Pair<Integer, Integer>>[] graph, int u, int prev,
                   List<Integer> leftBoundary, List<Integer> prefix,
                   Map<Integer, Integer> lastSeenDepth, int[] nums, int[] ans) {
    final int prevDepth = lastSeenDepth.getOrDefault(nums[u], 0);
    lastSeenDepth.put(nums[u], prefix.size());

    if (prevDepth != 0) {
      leftBoundary.add(prevDepth);
      Collections.sort(leftBoundary);
      leftBoundary = List.of(leftBoundary.get(leftBoundary.size() - 2),
                             leftBoundary.get(leftBoundary.size() - 1));
    }

    final int length = prefix.get(prefix.size() - 1) - prefix.get(leftBoundary.get(0));
    final int nodes = prefix.size() - leftBoundary.get(0);
    if (length > ans[0] || (length == ans[0] && nodes < ans[1])) {
      ans[0] = length;
      ans[1] = nodes;
    }

    for (Pair<Integer, Integer> pair : graph[u]) {
      final int v = pair.getKey();
      final int w = pair.getValue();
      if (v == prev)
        continue;
      prefix.add(prefix.get(prefix.size() - 1) + w);
      dfs(graph, v, u, new ArrayList<>(leftBoundary), prefix, lastSeenDepth, nums, ans);
      prefix.remove(prefix.size() - 1);
    }

    lastSeenDepth.put(nums[u], prevDepth);
  }
}

 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:
  # Similar to 3425. Longest Special Path
  def longestSpecialPath(
      self,
      edges: list[list[int]],
      nums: list[int]
  ) -> list[int]:
    maxLength = 0
    minNodes = 1
    graph = [[] for _ in range(len(nums))]

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

    prefix = [0]
    lastSeenDepth = {}

    def dfs(
        u: int,
        prev: int,
        leftBoundary: list[int],
    ) -> None:
      nonlocal maxLength, minNodes
      prevDepth = lastSeenDepth.get(nums[u], 0)
      lastSeenDepth[nums[u]] = len(prefix)

      if prevDepth != 0:
        leftBoundary = sorted(leftBoundary + [prevDepth])[-2:]

      length = prefix[-1] - prefix[leftBoundary[0]]
      nodes = len(prefix) - leftBoundary[0]
      if length > maxLength or (length == maxLength and nodes < minNodes):
        maxLength = length
        minNodes = nodes

      for v, w in graph[u]:
        if v == prev:
          continue
        prefix.append(prefix[-1] + w)
        dfs(v, u, leftBoundary)
        prefix.pop()

      lastSeenDepth[nums[u]] = prevDepth

    dfs(0, -1, leftBoundary=[0, 0])
    return [maxLength, minNodes]