LeetCode 1273: Delete Tree Nodes

https://leetcode.com › problems › delete-tree-nodes

Notes:

This is a good question.

Similar to the questions of binary search tree, we need to go through every sub-tree, to check the sum of it. If it equals to 0, all the node on that sub-tree can be deleted; otherwise, we need to return the number of nodes.

So the question becomes to calculate the sum and count the node number of every sub-tree. And this can be done by bottom-up recursion.

Since it is not a binary tree, we need to re-construct the tree first by using the information from parents and values.

The key information is: for each node (root of the sub-tree), we need to save all its children nodes.

After having this information, we can run a bottom-up recursion, and return two parameters: one is the sum, and the other is the node number.

See the code below:

class Solution {
public:
    int deleteTreeNodes(int nodes, vector<int>& parent, vector<int>& value) {
        int n = parent.size();
        map<int, vector<int>> mp;//<root, <kids>>
        for(int i=0; i<n; ++i) {
            mp[parent[i]].push_back(i);
        }
        auto res = dfs(-1, mp, value);
        return res.second;
    }

private:
    pair<int, int> dfs(int p, map<int, vector<int>>& mp, vector<int>& val) {
        int sum = p>=0 ? val[p] : 0, ct = p>=0 ? 1 : 0;
 if(mp.count(p)) {
            for(auto &a : mp[p]) {
                auto sub = dfs(a, mp, val);
                if(sub.first != 0) {
                    sum += sub.first;
                    ct += sub.second;
         }
            }
        }
        if(sum == 0) return {0, 0};
        return {sum, ct};
    }
};