NeetCodeのSolutionを書いていく

#

Kth Smallest Element In a Bst

#

問題

#

２分木のルートがと整数kが与えられます。k番目に小さい値を返しなさい。 左の木はノードのKeyより小さいKeyのノードしか含みません。 右の木はノードのKeyより大きいKeyのノードしか含みません。

例

#

Input: root = [2,1,3], k = 1 Output: 1

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

制約

#

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

メモ

#

• 小さい順に並べた時のK番目がほしい
• 左から順に調査すれば良い
• 再帰的に調査し、まずは一番左のノードを調査
• node.leftがNoneならば1を返す
• 次に右を調査、上記返却値を合わせて番目の数字を返すようにする
• 番目がKになった時のnode.valを返す うむ、どうしても複雑になってしまう。

答えを見ると、以下で行けるらしい

1. ノードをスタックしながら左に向かって調査
2. 左のターゲットがなくなったらスタックポップ+カウントアップ
3. この時点でカウントがkならvalを返す
4. そうではない場合は右の枝に現在地を移して エレガント！

from typing import Optional


# 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 kthSmallest(self, root: Optional[TreeNode], k: int) -> int:
        n = 0
        stack: list[TreeNode] = []
        cur = root

        while cur or stack:
            while cur:
                stack.append(cur)
                cur = cur.left

            cur = stack.pop()
            assert cur is not None

            n += 1
            if n == k:
                return cur.val
            cur = cur.right

        return n