NeetCode 150 [Trees]medium:Count Good Nodes In Binary Tree

NeetCodeのSolutionを書いていく

Count Good Nodes In Binary Tree

問題

バイナリツリーでnodexはrootからxまでのnodeでxよりも大きい値のnodeがない場合に良いノードとします。
バイナリツリーのrootが与えられた時、良いノードの数を返しなさい。

例

Input: root = [2,1,1,3,null,1,5]
Output: 3

Input: root = [2,1,1,3,null,1,5]
Output: 3

制約

• 1 <= number of nodes in the tree <= 100
• -100 <= Node.val <= 100
• 計算量：O(n)
• メモリ：O(n)

メモ

全ての親が自分より小さければ良い。
深さ優先探索で最大値を保持しながら評価すればいいか？

深さ優先探索を実現するには条件付きqueueを使えば良い？
更に、戻るときに同じノードを評価しないようにnodeのリンクと値を更新しておく。
また、nodeとそれまでの最大値を保持しておく

紹介されている回答は再帰で解いていてとてもスマート。

import collections


# Definition for a binary tree node.
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


class Solution:
    def goodNodes(self, root: TreeNode) -> int:
        q: collections.deque[tuple[TreeNode, int]] = collections.deque()
        q.append((root, root.val))
        maxv = root.val
        node = root
        # rootは必ず対象
        count = 1

        while True:
            # 右も左もない場合は上に戻る
            if (not node.left) and (not node.right):
                if len(q) == 0:
                    return count
                node, maxv = q.pop()
                continue

            # 左があれば左を次の探索対象にする
            temp = node
            if node.left:
                # 再評価されない様に更新
                node = node.left
                temp.left = None
                temp.val = -100
            elif node.right:
                node = node.right
                temp.right = None
                temp.val = -100

            # 新しいノードの値を評価
            if maxv <= node.val:
                maxv = node.val
                count += 1
            q.append((node, maxv))