Skip to content

1273. Delete Tree Nodes

  • Time: $O(n)$
  • Space: $O(h)$
 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
struct T {
  int sum;
  int count;
};

class Solution {
 public:
  int deleteTreeNodes(int nodes, vector<int>& parent, vector<int>& value) {
    vector<vector<int>> graph(nodes);

    for (int i = 1; i < parent.size(); ++i)
      graph[parent[i]].push_back(i);

    return dfs(graph, 0, value).count;
  }

 private:
  T dfs(const vector<vector<int>>& graph, int u, const vector<int>& value) {
    int sum = value[u];  // Root value
    int count = 1;       // This root

    for (const int v : graph[u]) {
      const T t = dfs(graph, v, value);
      sum += t.sum;
      count += t.count;
    }

    if (sum == 0)     // Delete this u
      return {0, 0};  // So its count = 0
    return {sum, count};
  }
};

 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
class T {
  public int sum;
  public int count;
  public T(int sum, int count) {
    this.sum = sum;
    this.count = count;
  }
}

class Solution {
  public int deleteTreeNodes(int nodes, int[] parent, int[] value) {
    List<Integer>[] graph = new List[nodes];

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

    for (int i = 1; i < parent.length; ++i)
      graph[parent[i]].add(i);

    return dfs(graph, 0, value).count;
  }

  private T dfs(List<Integer>[] graph, int u, int[] value) {
    int sum = value[u]; // Root value
    int count = 1;      // This root

    for (final int v : graph[u]) {
      final T t = dfs(graph, v, value);
      sum += t.sum;
      count += t.count;
    }

    if (sum == 0)         // Delete this u
      return new T(0, 0); // So its count = 0
    return new T(sum, count);
  }
}