ลบโหนดต้นไม้ใน C++

สมมติว่าเรามีต้นไม้ที่รากอยู่ที่โหนด 0 กำหนดโดย:

1. จำนวนโหนดคือ nodes
2. ค่าของโหนดที่ i คือ value[i]
3. พ่อของโหนดที่ i คือ parent[i]

เราต้องลบ subtree ทุกตัวที่ผลรวมค่าของโหนดเป็น 0 แล้วคืนจำนวนโหนดที่เหลือ

ตัวอย่าง ต้นไม้ดังภาพ มี 7 โหนด ผลลัพธ์: 2

ขั้นตอนการแก้ปัญหา (DFS Post-order)

• สร้างแผนที่ children
• กำหนดเมธอด dfs(node, value, graph) คืนคู่ (sum, count):
  • temp := (value[node], 1)
  • วนลูปลูกทุกตัวของ node:
    • temp2 := dfs(child, value, graph)
    • บวก temp.first += temp2.first, temp.second += temp2.second
  • ถ้า temp.first == 0:
    • ans := ans - temp.second
    • temp.second := 0
  • คืน temp
• ในเมธอดหลัก:
  • n := ขนาดของ value, ans := n
  • สร้างกราฟ graph ขนาด n+1
  • วนลูป i จาก 1 ถึง n-1 ใส่ i เข้า graph[parent[i]]
  • เรียก dfs(0, value, graph)
  • คืน ans

ตัวอย่างโค้ด (C++)

#include <bits/stdc++.h> 
using namespace std; 

class Solution {
public:
    map <int, int> children;
    int ans;
    pair <int, int> dfs(int node, vector<int>& value, vector <int> graph[]) {
        pair <int, int> temp = {value[node], 1};
        for(int i = 0; i < graph[node].size(); i++) {
            pair <int, int> temp2 = dfs(graph[node][i], value, graph);
            temp.first += temp2.first;
            temp.second += temp2.second;
        }
        if(temp.first == 0) {
            ans -= temp.second;
            temp.second = 0;
        }
        return temp;
    }
    int deleteTreeNodes(int nodes, vector<int>& parent, vector<int>& value) {
        int n = value.size();
        ans = n;
        children.clear();
        vector < int > graph[n + 1];
        for(int i = 1; i < n; i++) {
            graph[parent[i]].push_back(i);
        }
        dfs(0, value, graph);
        return ans;
    }
};

int main() {
    vector<int> v1 = {-1,0,0,1,2,2,2};
    vector<int> v2 = {1,-2,4,0,-2,-1,-1};
    Solution ob;
    cout << (ob.deleteTreeNodes(7, v1, v2));
}

อินพุต

7 [-1,0,0,1,2,2,2] [1,-2,4,0,-2,-1,-1]

เอาต์พุต

2