【木のルールどこだ】高さ平衡二分探索木問題を解いてみた

解いた問題

leetcode 108. Convert Sorted Array to Binary Search Tree
ソートされた配列を高さ平衡二分探索木に変換してねって問題。

高さ平衡二分探索木とは、各ノードにおいて、「左の部分期にあるすべての値 < 自分の値 < 右の部分期にあるすべての値」が成り立つ木構造のこと。

最初に考えた解法

miss.py

class Solution:
    def sortedArrayToBST(self, nums: List[int]) -> Optional[TreeNode]:
        results = []
        size = len(nums)
        mid = size // 2 
        val = nums[mid]
        left = nums[:mid]
        right = nums[mid+1:]
        result += mid
        sortedArrayToBST(left)
        sortedArrayToBST(right)

ルートと左右ノードに分け、ルートから返却するリストに突っ込んでいく、左右ノードであるleftとrightを再帰処理でまたルートと左右ノードに分けて返却リストに大小考えて突っ込む、ということを考えた。

途中まで考えてmidのサイズが1になった時インデックスエラーが出そうなことに気づき、そのハンドリングを組み込もうとすると冗長になると感じギブアップ。

回答

answer.py

# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def sortedArrayToBST(self, nums: List[int]) -> Optional[TreeNode]:
        if not nums:
            return None
        mid = len(nums) // 2

        root = TreeNode(nums[mid])
        root.left = self.sortedArrayToBST(nums[:mid])
        root.right = self.sortedArrayToBST(nums[mid+1:])

        return root

「ルート、左ノード、右ノード」という３つのノードの塊を一つのインスタンスとしてとらえ、ノードが次のインスタンスへのポインタを保持することで木構造を作るという処理になっている。

感想

考え自体はあっていたのと、左右への分け方が同じでうれしかった。
認識がずれていたところは「戻り値の型」について。

てっきりlistを返却するものかと思ったらOptional[TreeNode]ってガッツリ書いてあった。
「TreeNodeインスタンスを返却してくれたら内部的に勝手に成形しますよ」ってことね。

また、TreeNodeクラス内において、デフォルト引数をNoneにすることによって勝手に枝分かれする要素がない場合も定義できているのかと感心。

あと「木構造とは左にある数字が小さくて右に行くほど大きいっすよ。ルート部分含めてそうっすよ」ってことはどこにも書かれていないのね。
初心者には優しくないと思いつつ、そんなこと書いてたら冗長で見づらくなりそうとも感じた。