LeetCode / Balanced Binary Tree

https://leetcode.com/problems/balanced-binary-tree/

Given a binary tree, determine if it is height-balanced.

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 1:

Given the following tree [3,9,20,null,null,15,7]:

Return true.

Example 2:

Given the following tree [1,2,2,3,3,null,null,4,4]:

Return false.

どのsubtree間の深さの違いが1以下であるtreeを、balanced binary treeと言っています。
treeがbalanced binary treeになっているか判定せよという問題です。

解答・解説

解法1

subtreeの深さの違いからbalancedな状態か判定する問題なので、subtreeの深さを取得しながら、左右のsubtreeの深さの違いが1以下であるか判定する、再帰的な処理を書くこととします。
以下ではcheck関数を定義し、(subtreeの深さ, 左右のsubtree間の深さの差が1以下であるかのboolean)をtupleで返す構造とし、rootがない末端では(0, True）を返すこととして、再帰的な処理をかけます。
isBalanced関数の返り値としては、tupleの2つ目の要素を返すことになります。

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

def check(root):
    if not root:
        return (0, True)
    l_depth, l_balanced = check(root.left)
    r_depth, r_balanced = check(root.right)
    return max(l_depth, r_depth) + 1, l_balanced and r_balanced and abs(l_depth - r_depth) <= 1

class Solution:
    def isBalanced(self, root):
        return check(root)[1]

iterativeなコードはちょっと長くなるので、recursiveが良いと思います。
考え方は同じですが、違う書き方をすると例えば以下の通り。

def check(root):
    if not root: return 0
    l_depth = check(root.left)
    r_depth = check(root.right)
    if l_depth == -1 or r_depth == -1 or abs(l_depth - r_depth) > 1: return -1
    return max(l_depth, r_depth) + 1

class Solution:
    def isBalanced(self, root):
        return check(root) != -1