1. 程式人生 > >[LeetCode] Equal Tree Partition 劃分等價樹

[LeetCode] 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.

這道題讓我們劃分等價樹,就是說當移除一條邊後,被分成的兩棵樹的結點之和需要相等。那麼通過觀察題目中的例子我們可以發現,如果我們將每個結點的結點值變成其所有子結點的結點值之和再加上當前的結點值,那麼對於例子1來說,根結點的結點值就變成了30,斷開位置的結點就變成了15,那麼我們就可以發現其實只要斷開位置的結點值是根結點值的一半,就存在等價劃分。所以這道題的難點就是更新每個結點的結點值,我們可以使用遞迴來做。博主最開始使用的是unordered_set,把更新後的每個結點值都存入集合中,但是對於test case: [0, 1, -1] 會fail, 仔細分析下這個case,發現更新後的根結點值還是0,而且0已經被存入集合了,而0除以2還是0,在集合中存在,會返回true,但其實這棵樹是不能等價劃分的。0的情況確實比較特殊,所以我們使用unordered_map,建立更新後的結點值和其出現次數之間的對映,這樣只有map中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; } };

參考資料: