统计和生成所有不同的二叉树

统计和生成所有不同的二叉树

题目

给定一个整数n，如果n<1，代表空树，否则代表中序遍历的结果为1,2,3…n，求可能的二叉树结构有多少。例如：
n=-1表示空树，返回1
n=2，中序遍历的1，2，满足的二叉树有：

思路

假设f(n)表示中序遍历序列为1,2…n可表示的二叉树数目，该树一定是二叉搜索树。
则1,2,…n都可以表示根节点，假设1<=a<=n，a为根的二叉搜索树种数： f(a-1) * f(n-a)
则f(n) = f(i-1)*f(n-a), for i in [1,n]

实现1，递归

@timethis
def num_bst1(n):
    def recursive(n):
        if n < 2:
            return 1

        num = 0
        for i in range(1, n+1):
            num += recursive(i-1) * recursive(n-i)

        return num

    return recursive(n)

包含大量重复计算，可以预计这个实现效率不会太高

实现2，带记忆的递归

@timethis
def num_bst2(n):
    def recursive(n):
        if n in d:
            return d[n]

        if n < 2:
            d[n] = 1
            return 1

        num = 0
        for i in range(1, n+1):
            num += recursive(i-1) * recursive(n-i)

        d[n] = num
        return num

    d = {}
    return recursive(n)

实现3，动态规划

计算f(n)依赖f(1), f(2) … f(i)，这种是动态规划发挥作用的好地方：

@timethis
def num_bst(n):
    if n < 2:
        return 1

    dp = [0 for i in range(n+1)]
    dp[0] = 1

    for i in range(1, n+1):
        for j in range(1, i+1):
            dp[i] += dp[j-1] * dp[i-j]

    return dp[n]

测试

import time
from functools import wraps


class TreeNode():
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None


def timethis(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print(func.__name__, 'cost time:', end - start)
        return result
    return wrapper

def test(n):
    n0 = num_bst(n)
    n1 = num_bst1(n)
    n2 = num_bst2(n)
    print(n0, n1, n2)

if __name__ == '__main__':
    test(3)
    test(8)
    test(13)
    test(33)

结果

num_bst cost time: 6.4373016357421875e-06
num_bst1 cost time: 5.0067901611328125e-06
num_bst2 cost time: 5.0067901611328125e-06
5 5 5
num_bst cost time: 9.5367431640625e-06
num_bst1 cost time: 0.0005745887756347656
num_bst2 cost time: 1.3589859008789062e-05
1430 1430 1430
num_bst cost time: 1.5735626220703125e-05
num_bst1 cost time: 0.16288375854492188
num_bst2 cost time: 4.2438507080078125e-05
742900 742900 742900
num_bst cost time: 0.0001049041748046875
^CTraceback (most recent call last):
  File "count_and_generate_bst.py", line 147, in <module>
    test(33)
  File "count_and_generate_bst.py", line 132, in test
    n1 = num_bst1(n)
  File "count_and_generate_bst.py", line 16, in wrapper
    result = func(*args, **kwargs)
  File "count_and_generate_bst.py", line 50, in num_bst1
    return recursive(n)
  File "count_and_generate_bst.py", line 46, in recursive
    num += recursive(i-1) * recursive(n-i)
  File "count_and_generate_bst.py", line 46, in recursive
    num += recursive(i-1) * recursive(n-i)
  File "count_and_generate_bst.py", line 46, in recursive
    num += recursive(i-1) * recursive(n-i)
  [Previous line repeated 21 more times]
  File "count_and_generate_bst.py", line 45, in recursive
    for i in range(1, n+1):
KeyboardInterrupt

n=33时，单纯的递归已经好几秒算不出来了，去掉它，只用另外两种算法：

num_bst cost time: 5.245208740234375e-06
num_bst2 cost time: 4.5299530029296875e-06
5 5
num_bst cost time: 1.8358230590820312e-05
num_bst2 cost time: 3.123283386230469e-05
1430 1430
num_bst cost time: 1.4543533325195312e-05
num_bst2 cost time: 5.936622619628906e-05
742900 742900
num_bst cost time: 0.00011610984802246094
num_bst2 cost time: 0.0001766681671142578
212336130412243110 212336130412243110

进阶

生成这些树，返回生成的这些树的根节点。

def generate(n):
    def clone_tree(root):
        if root is None:
            return None

        r = TreeNode(root.val)
        r.left = clone_tree(root.left)
        r.right = clone_tree(root.right)

        return r

    def recursive(start, end):
        res = []
        if start > end:
            return [None]

        root = None
        for i in range(start, end+1):
            root = TreeNode(i)
            left_subs = recursive(start, i-1)
            right_subs = recursive(i+1, end)
            for l in left_subs:
                for r in right_subs:
                    root.left, root.right = l, r
                    res.append(clone_tree(root))

        return res

    return recursive(1, n)

测试

def print_tree_by_level(root):
    if root is None:
        return

    q = [root, None]
    level = 0
    print('level %d: ' % (level), end='')
    while q:
        node = q.pop(0)
        if node is None:
            if q:
                level += 1
                q.append(None)
                print()
                print('level %d: ' % (level), end='')
        else:
            print(node.val, end=', ')
            if node.left:
                q.append(node.left)
            if node.right:
                q.append(node.right)
    print()
    print('------------------------')


def test(n):
    n0 = num_bst(n)
    #n1 = num_bst1(n)
    n2 = num_bst2(n)
    #print(n0, n1, n2)
    print(n0, n2)


def test_generate(n):
    results = generate(n)
    for root in results:
        print_tree_by_level(root)


if __name__ == '__main__':
    test(3)
    test(8)
    test(13)
    test(33)
    test_generate(3)

结果

num_bst cost time: 5.245208740234375e-06
num_bst2 cost time: 4.5299530029296875e-06
5 5
num_bst cost time: 1.8358230590820312e-05
num_bst2 cost time: 3.123283386230469e-05
1430 1430
num_bst cost time: 1.4543533325195312e-05
num_bst2 cost time: 5.936622619628906e-05
742900 742900
num_bst cost time: 0.00011610984802246094
num_bst2 cost time: 0.0001766681671142578
212336130412243110 212336130412243110
level 0: 1, 
level 1: 2, 
level 2: 3, 
------------------------
level 0: 1, 
level 1: 3, 
level 2: 2, 
------------------------
level 0: 2, 
level 1: 1, 3, 
------------------------
level 0: 3, 
level 1: 1, 
level 2: 2, 
------------------------
level 0: 3, 
level 1: 2, 
level 2: 1, 
------------------------