2467. Most Profitable Path in a Tree

• Time: $O(n)$
• Space: $O(n)$

C++JavaPython

class Solution {
 public:
  int mostProfitablePath(vector<vector<int>>& edges, int bob,
                         vector<int>& amount) {
    const int n = amount.size();
    vector<vector<int>> tree(n);
    vector<int> parent(n);
    vector<int> aliceDist(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);
    }

    dfs(tree, 0, -1, 0, parent, aliceDist);

    // Modify the amount along the path from node Bob to node 0.
    // For each node,
    //   1. If Bob reaches earlier than Alice does, change the amount to 0.
    //   2. If Bob and Alice reach simultaneously, devide the amount by 2.
    for (int u = bob, bobDist = 0; u != 0; u = parent[u], ++bobDist)
      if (bobDist < aliceDist[u])
        amount[u] = 0;
      else if (bobDist == aliceDist[u])
        amount[u] /= 2;

    return getMoney(tree, 0, -1, amount);
  }

 private:
  // Fills `parent` and `dist`.
  void dfs(const vector<vector<int>>& tree, int u, int prev, int d,
           vector<int>& parent, vector<int>& dist) {
    parent[u] = prev;
    dist[u] = d;
    for (const int v : tree[u]) {
      if (dist[v] == -1)
        dfs(tree, v, u, d + 1, parent, dist);
    }
  }

  int getMoney(const vector<vector<int>>& tree, int u, int prev,
               const vector<int>& amount) {
    // a leaf node
    if (tree[u].size() == 1 && tree[u][0] == prev)
      return amount[u];

    int maxPath = INT_MIN;
    for (const int v : tree[u])
      if (v != prev)
        maxPath = max(maxPath, getMoney(tree, v, u, amount));

    return amount[u] + maxPath;
  }
};

class Solution {
  public int mostProfitablePath(int[][] edges, int bob, int[] amount) {
    final int n = amount.length;
    List<Integer>[] tree = new List[n];
    int[] parent = new int[n];
    int[] aliceDist = new int[n];
    Arrays.fill(aliceDist, -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);
    }

    dfs(tree, 0, -1, 0, parent, aliceDist);

    // Modify the amount along the path from node Bob to node 0.
    // For each node,
    //   1. If Bob reaches earlier than Alice does, change the amount to 0.
    //   2. If Bob and Alice reach simultaneously, devide the amount by 2.
    for (int u = bob, bobDist = 0; u != 0; u = parent[u], ++bobDist)
      if (bobDist < aliceDist[u])
        amount[u] = 0;
      else if (bobDist == aliceDist[u])
        amount[u] /= 2;

    return getMoney(tree, 0, -1, amount);
  }

  // Fills `parent` and `dist`.
  private void dfs(List<Integer>[] tree, int u, int prev, int d, int[] parent, int[] dist) {
    parent[u] = prev;
    dist[u] = d;
    for (final int v : tree[u]) {
      if (dist[v] == -1)
        dfs(tree, v, u, d + 1, parent, dist);
    }
  }

  private int getMoney(List<Integer>[] tree, int u, int prev, int[] amount) {
    // a leaf node
    if (tree[u].size() == 1 && tree[u].get(0) == prev)
      return amount[u];

    int maxPath = Integer.MIN_VALUE;
    for (final int v : tree[u])
      if (v != prev)
        maxPath = Math.max(maxPath, getMoney(tree, v, u, amount));

    return amount[u] + maxPath;
  }
}

class Solution:
  def mostProfitablePath(
      self,
      edges: list[list[int]],
      bob: int,
      amount: list[int],
  ) -> int:
    n = len(amount)
    tree = [[] for _ in range(n)]
    parent = [0] * n
    aliceDist = [-1] * n

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

    # Fills `parent` and `aliceDist`.
    def dfs(u: int, prev: int, d: int) -> None:
      parent[u] = prev
      aliceDist[u] = d
      for v in tree[u]:
        if aliceDist[v] == -1:
          dfs(v, u, d + 1)

    dfs(0, -1, 0)

    # Modify amount athe path from node bob to node 0.
    # For each node,
    #   1. If Bob reaches earlier than Alice does, change the amount to 0.
    #   2. If Bob and Alice reach simultaneously, devide the amount by 2.
    u = bob
    bobDist = 0
    while u != 0:
      if bobDist < aliceDist[u]:
        amount[u] = 0
      elif bobDist == aliceDist[u]:
        amount[u] //= 2
      u = parent[u]
      bobDist += 1

    return self._getMoney(tree, 0, -1, amount)

  def _getMoney(
      self,
      tree: list[list[int]],
      u: int,
      prev: int,
      amount: list[int],
  ) -> int:
    # a leaf node
    if len(tree[u]) == 1 and tree[u][0] == prev:
      return amount[u]

    maxPath = -math.inf
    for v in tree[u]:
      if v != prev:
        maxPath = max(maxPath, self._getMoney(tree, v, u, amount))

    return amount[u] + maxPath