使用python建立简单的树机构

代码

import sys


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


class Solution:
    def preorderTraversal(self, root):
        """
        :type root: TreeNode
        :rtype: List[int]
        """
        ret = []
        stack = [root]
        while stack:
            node = stack.pop()
            if node:
                ret.append(node.val)
                stack.append(node.right)
                stack.append(node.left)
        return ret


def list_to_treenode(input_values):
    if not input_values:
        return None
    root = TreeNode(int(input_values[0]))
    node_queue = [root]
    front = 0
    index = 1
    while index < len(input_values):
        node = node_queue[front]
        front += 1

        item = input_values[index]
        index += 1

        if item != "null":
            left_num = int(item)
            node.left = TreeNode(left_num)
            node_queue.append(node.left)

        if index >= len(input_values):
            break

        item = input_values[index]
        index += 1
        if item != "null":
            right_num = int(item)
            node.right = TreeNode(right_num)
            node_queue.append(node.right)
    return root


def main():
    line = sys.stdin.readline().strip()
    values = list(map(str, line.split()))
    root = list_to_treenode(values)
    pret = Solution().preorderTraversal(root)
    print(pret)


if __name__ == '__main__':
    main()

测试

# 输入
3 9 20 null null 15 7

    3
   / \
  9  20
    /  \
   15   7

# 前序遍历
[3, 9, 20, 15, 7]