1273. Delete Tree Nodes

A tree rooted at node 0 is given as follows:

  • The number of nodes is nodes;
  • The value of the i-th node is value[i];
  • The parent of the i-th node is parent[i].

Remove every subtree whose sum of values of nodes is zero.

After doing so, return the number of nodes remaining in the tree.

 

Example 1:

Input: nodes = 7, parent = [-1,0,0,1,2,2,2], value = [1,-2,4,0,-2,-1,-1]
Output: 2

Example 2:

Input: nodes = 7, parent = [-1,0,0,1,2,2,2], value = [1,-2,4,0,-2,-1,-2]
Output: 6

Example 3:

Input: nodes = 5, parent = [-1,0,1,0,0], value = [-672,441,18,728,378]
Output: 5

Example 4:

Input: nodes = 5, parent = [-1,0,0,1,1], value = [-686,-842,616,-739,-746]
Output: 5

 

Constraints:

  • 1 <= nodes <= 10^4
  • parent.length == nodes
  • 0 <= parent[i] <= nodes - 1
  • parent[0] == -1 which indicates that 0 is the root.
  • value.length == nodes
  • -10^5 <= value[i] <= 10^5
  • The given input is guaranteed to represent a valid tree.

1273. Delete Tree Nodes
struct Solution;
use std::collections::HashMap;

impl Solution {
    fn delete_tree_nodes(nodes: i32, parent: Vec<i32>, value: Vec<i32>) -> i32 {
        let mut hm: HashMap<usize, Vec<usize>> = HashMap::new();
        let n = nodes as usize;
        let mut root = n;
        for i in 0..n {
            if parent[i] != -1 {
                hm.entry(parent[i] as usize).or_default().push(i);
            } else {
                root = i;
            }
        }
        let (_, size) = Self::postorder(root, &hm, &value, n);
        size as i32
    }

    fn postorder(
        i: usize,
        hm: &HashMap<usize, Vec<usize>>,
        value: &[i32],
        n: usize,
    ) -> (i32, usize) {
        let mut sum = value[i];
        let mut size = 1;
        if let Some(children) = hm.get(&i) {
            for &j in children {
                let child = Self::postorder(j, hm, value, n);
                sum += child.0;
                size += child.1;
            }
        }
        if sum == 0 {
            (0, 0)
        } else {
            (sum, size)
        }
    }
}

#[test]
fn test() {
    let nodes = 7;
    let parent = vec![-1, 0, 0, 1, 2, 2, 2];
    let value = vec![1, -2, 4, 0, -2, -1, -1];
    let res = 2;
    assert_eq!(Solution::delete_tree_nodes(nodes, parent, value), res);
}