NeetCode 150 [Trees]medium:Kth Smallest Element In a Bst

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