นับ Subtree ที่ผลรวมน้ำหนักเท่ากับ x ใน Binary Tree ด้วย C++

โจทย์: ให้ Binary Tree และค่า x ต้องการนับจำนวน Subtree ที่ผลรวมน้ำหนักโหนดเท่ากับ x

แนวคิด: Post-order Traversal คำนวณผลรวม

คำนวณผลรวมของ Subtree ทุกตัวจากใบขึ้นมา หากผลรวมเท่ากับ x ให้นับ +1

ฟังก์ชันหลัก

int subtrees_x(Tree_Node* root, int x) {
    static int count = 0;
    static Tree_Node* temp = root;
    int Left_subtree = 0, Right_subtree = 0;
    
    if (root == NULL) return 0;
    
    Left_subtree += subtrees_x(root->Left, x);
    Right_subtree += subtrees_x(root->Right, x);
    
    int sum = Left_subtree + Right_subtree + root->data;
    
    if (sum == x) count++;
    
    if (temp != root)  // ไม่ใช่รากต้นไม้หลัก
        return Left_subtree + root->data + Right_subtree;
    
    return count;
}

ตัวอย่างเต็ม

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

struct Tree_Node {
    int data;
    Tree_Node *Left, *Right;
};

Tree_Node* insert_Node(int data) {
    Tree_Node* new_node = (Tree_Node*)malloc(sizeof(Tree_Node));
    new_node->data = data;
    new_node->Left = new_node->Right = NULL;
    return new_node;
}

int subtrees_x(Tree_Node* root, int x) {
    static int count = 0;
    static Tree_Node* temp = root;
    int Left_subtree = 0, Right_subtree = 0;
    if (root == NULL) return 0;
    Left_subtree += subtrees_x(root->Left, x);
    Right_subtree += subtrees_x(root->Right, x);
    int sum = Left_subtree + Right_subtree + root->data;
    if (sum == x) count++;
    if (temp != root) {
        int set = Left_subtree + root->data + Right_subtree;
        return set;
    }
    return count;
}

int main() {
    Tree_Node* root = insert_Node(10);
    root->Left = insert_Node(20);
    root->Right = insert_Node(12);
    root->Left->Left = insert_Node(14);
    root->Left->Right = insert_Node(1);
    root->Right->Left = insert_Node(21);
    root->Right->Right = insert_Node(11);
    int x = 14;
    cout << "Count: " << subtrees_x(root, x);
    return 0;
}

ผลลัพธ์

Count of subtrees that sum up to a given value x are: 1

คำอธิบาย: มีแค่โหนดใบค่า 14 ที่มีผลรวมเท่ากับ 14 (ตัวมันเอง)

ความซับซ้อน: Time O(n), Space O(h)