64 lines
1.5 KiB
Python
64 lines
1.5 KiB
Python
# -*- coding: UTF-8 -*-
|
|
|
|
|
|
class TreeNode(object):
|
|
def __init__(self, x):
|
|
self.val = x
|
|
self.left = None
|
|
self.right = None
|
|
|
|
def __repr__(self):
|
|
return 'node %s' % self.val
|
|
|
|
|
|
def creatBTree(data, index):
|
|
"""通过数组生成二叉树
|
|
|
|
Args:
|
|
data ([int, None]): 每个节点都要定义，包括空节点 eg. [1,None,2, None, None, 3]
|
|
index (int): 返回哪个位置的节点
|
|
|
|
Returns:
|
|
TreeNode: TreeNode obj
|
|
"""
|
|
pNode = None
|
|
if index < len(data):
|
|
if data[index] == None:
|
|
return
|
|
pNode = TreeNode(data[index])
|
|
pNode.left = creatBTree(data, 2 * index + 1) # [1, 3, 7, 15, ...]
|
|
pNode.right = creatBTree(data, 2 * index + 2) # [2, 5, 12, 25, ...]
|
|
return pNode
|
|
|
|
|
|
def generate_tree(vals):
|
|
"""通过 leetcode 样式的数组生成二叉树
|
|
|
|
Args:
|
|
vals ([int, None]): leetcode 格式，省略全空子节点 eg. [1, None, 2, 3]
|
|
|
|
Returns:
|
|
TreeNode: TreeNode obj
|
|
"""
|
|
if not vals:
|
|
return None
|
|
que = []
|
|
fill_left = True
|
|
for val in vals:
|
|
node = TreeNode(val) if val else None
|
|
if len(que) == 0:
|
|
root = node
|
|
que.append(node)
|
|
elif fill_left:
|
|
que[0].left = node
|
|
fill_left = False
|
|
if node:
|
|
que.append(node)
|
|
else:
|
|
que[0].right = node
|
|
if node:
|
|
que.append(node)
|
|
que.pop(0)
|
|
fill_left = True
|
|
return root
|