DFS應用題 計算與子樹平均值相等的節點數 Leetcode #2265

題目敘述

題目會給一顆二元樹，要求我們計算節點值 和 子樹平均值相等的node有幾個。

詳細的題目敘述在這裡:

測試範例

Example 1:

Input: root = [4,8,5,0,1,null,6]
Output: 5
Explanation: 
For the node with value 4: The average of its subtree is (4 + 8 + 5 + 0 + 1 + 6) / 6 = 24 / 6 = 4.
For the node with value 5: The average of its subtree is (5 + 6) / 2 = 11 / 2 = 5.
For the node with value 0: The average of its subtree is 0 / 1 = 0.
For the node with value 1: The average of its subtree is 1 / 1 = 1.
For the node with value 6: The average of its subtree is 6 / 1 = 6.

Example 2:

Input: root = [1]
Output: 1
Explanation: For the node with value 1: The average of its subtree is 1 / 1 = 1.

約束條件

Constraints:

• The number of nodes in the tree is in the range [1, 1000].
• 0 <= Node.val <= 1000

演算法

抽象來想，對於當下節點來說，需要知道下列資訊來滿足題目要求的計算

1. 子樹的結點值總和
2. 子樹的valid count (就是子樹內部，節點值=子樹平均值的node count)
3. 子樹的總結點數(因為平均 = 子樹結點值總和/子樹總結點數，我們會需要分子和分母的確切數值)

依循這個思路，設計DFS 遞迴function，以bottom-up 由下往上的方式，回傳 上述三個參數給parent node，逐層計算節點值 和 子樹平均值相等的node有幾個。

程式碼

class Solution:
　def averageOfSubtree(self, root: Optional[TreeNode]) -> int:
　　
　　def helper(root):

　　　if not root:
　　　　# node value, valid count, total node count
　　　　return 0, 0, 0
　　　
　　　l_sum, l_valid, l_nodes = helper(root.left)

　　　r_sum, r_valid, r_nodes = helper(root.right)

　　　# Average = (root + left subtree + right subtree) / subtree node count
　　　avg = (root.val + l_sum + r_sum) // (1 + l_nodes + r_nodes )

　　　if root.val == avg:
　　　　return l_sum + r_sum + root.val, l_valid + r_valid + 1, l_nodes + r_nodes + 1
　　　
　　　else:
　　　　return l_sum + r_sum + root.val, l_valid + r_valid, l_nodes + r_nodes + 1

　　# ===========================
　　# Our goal is valid count
　　return helper(root)[1]

複雜度分析

時間複雜度:

O( n ) DFS 拜訪每個節點，每個節點拜訪一次。

空間複雜度:

O( n ) DFS 拜訪每個節點，最深深度為O(n)，發生在整棵樹向左歪斜或向右歪斜時。

關鍵知識點

DFS function 傳統上都是回傳一個參數，或者直接return不帶回傳值。
這題的關鍵考點在於，我們需要的各種資訊都隱含在子樹裡，而且不只一個欄位。
接著想到，function其實可以return 多重欄位的回傳值，回傳給parent node 去計算題目所要求的目標值。

Reference:

[1] Python O(n) by DFS - Count Nodes Equal to Average of Subtree - LeetCode