オットセイの経営日誌

データサイエンス系ベンチャーを経営してます。経営のこと、趣味のことつぶやきます。

LeetCode / Convert Sorted Array to Binary Search Tree

https://leetcode.com/problems/convert-sorted-array-to-binary-search-tree/

Given an array where elements are sorted in ascending order, convert it to a height balanced BST.

For this problem, a height-balanced binary tree is defined as a binary tree in which the depth of the two subtrees of every node never differ by more than 1.

Example:

f:id:mhiro216:20190730083844p:plain:w600

a height-balanced binary tree is defined as a binary tree in which the depth of the two subtrees of every node never differ by more than 1ということで、
2つのsubtree間の深さに2以上の違いがないtree」を作れと言われています。

解答・解説

解法1

2つのsubtree間の深さも1より大きな違いがないようなtreeは色々考えられるわけですが、問題文にもあった下図のように、

f:id:mhiro216:20190730203659p:plain:w100 (binarytreeライブラリで作ってみましたがイマイチですね。。。)

大元のrootから左右に2本subtreeが出て、あとは1本しかsubtreeがないようなtreeを考え、2つのsubtreeの深さに2以上の違いがないtree
を作るのがシンプルなアルゴリズムになります。より具体的には、
[-10, -3, 0, 5, 9] -> left: [-10, -3], root: 0, right: [5, 9]
のように、中央値をrootとし、中央値より前の値のリストをroot.leftに、中央値より後の値のリストをroot.rightとします。
そして、root.leftとroot.rightに対して同じ処理をする再帰的な構造にします。

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

class Solution: 
    def sortedArrayToBST(self, nums: List[int]) -> 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
解法2

Discussionを眺めていたら、リストのスライスはコストが高いということで、修正案がありました。

class Solution: 
    def sortedArrayToBST(self, nums: List[int]) -> TreeNode:
        def convert(left, right):
            if left > right:
                return None
            mid = (left + right) // 2
            node = TreeNode(nums[mid])
            node.left = convert(left, mid - 1)
            node.right = convert(mid + 1, right)
            return node
        return convert(0, len(nums) - 1)

Python操作に関するコストは以下のページにまとまっています。

wiki.python.org