[LeetCode] 663. Equal Tree Partition 划分等價樹

 

Given a binary tree with n nodes, your task is to check if it's possible to partition the tree to two trees which have the equal sum of values after removing exactly one edge on the original tree.

Example 1:

Input:     
    5
   / \
  10 10
    /  \
   2   3

Output: True
Explanation: 
    5
   / 
  10
      
Sum: 15

   10
  /  \
 2    3

Sum: 15

 

Example 2:

Input:     
    1
   / \
  2  10
    /  \
   2   20

Output: False
Explanation: You can't split the tree into two trees with equal sum after removing exactly one edge on the tree.

 

Note:

1. The range of tree node value is in the range of [-100000, 100000].
2. 1 <= n <= 10000

 

這道題讓我們划分等價樹，就是說當移除一條邊后，被分成的兩棵樹的結點之和需要相等。通過觀察題目中的例子可以發現，如果將每個結點的結點值變成其所有子結點的結點值之和再加上當前的結點值，那么對於例子1來說，根結點的結點值就變成了 30，斷開位置的結點就變成了 15，可以發現其實只要斷開位置的結點值是根結點值的一半，就存在等價划分。所以這道題的難點就是更新每個結點的結點值，可以使用遞歸來做。博主最開始使用的是 unordered_set，把更新后的每個結點值都存入集合中，但是對於 test case: [0, 1, -1] 會 fail, 仔細分析下這個 case，發現更新后的根結點值還是0，而且0已經被存入集合了，而0除以2還是0，在集合中存在，會返回 true，但其實這棵樹是不能等價划分的。0的情況確實比較特殊，所以這里要使用 HashMap，建立更新后的結點值和其出現次數之間的映射，這樣只有 HashMap 中0的個數大於1的時候，才返回 true。這樣完美的避開了根結點為0的陷阱，Perfect！參見代碼如下：

 

解法一：

class Solution {
public:
    bool checkEqualTree(TreeNode* root) {
        unordered_map<int, int> m;
        int sum = helper(root, m);
        if (sum == 0) return m[0] > 1;
        return (sum % 2 == 0) && m.count(sum / 2);
    }
    int helper(TreeNode* node, unordered_map<int, int>& m) {
        if (!node) return 0;
        int cur = node->val + helper(node->left, m) + helper(node->right, m);
        ++m[cur];
        return cur;
    }
};

 

我們也可以使用 stack 來做，將所有的結點和存入到棧中，然后依次出棧，看是否有結點和正好等於 sum/2，這里還是要注意當 sum=0 的情況同時根結點值又正好是0的情況，最簡單的處理辦法就是將棧頂元素提前移除，這樣就不會有 sum=sum/2 這種情況發生了，參見代碼如下：

 

解法二：

class Solution {
public:
    bool checkEqualTree(TreeNode* root) {
        stack<int> st;
        int sum = helper(root, st);
        st.pop();
        if (sum % 2 != 0) return false;
        while (!st.empty()) {
            if (st.top() == sum / 2) return true;
            st.pop();
        }
        return false;
    }
    int helper(TreeNode* node, stack<int>& st) {
        if (!node) return 0;
        st.push(helper(node->left, st) + helper(node->right, st) + node->val);
        return st.top();
    }
};

 

Github 同步地址：

https://github.com/grandyang/leetcode/issues/663

 

參考資料：

https://leetcode.com/problems/equal-tree-partition/

https://leetcode.com/problems/equal-tree-partition/discuss/149437/Logical-Thinking-with-Java-Code-Beats-97.25

https://leetcode.com/problems/equal-tree-partition/discuss/106727/JavaC%2B%2B-Simple-solution-with-only-one-HashMaplessgreater.

 

LeetCode All in One 題目講解匯總(持續更新中...)