在二叉树中找到一个节点的后继节点

在二叉树中找到一个节点的后继节点

题目

假设二叉搜索树节点带有父节点指针：

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

给定二叉搜索树上一个节点node，返回node的后继节点

BST

class BST():
    def __init__(self, root=None):
        self.root = root

    def insert_recursive(self, val):
        def recursive(node, val):
            if not node:
                return TreeNode(val)

            if val < node.val:
                node.left = recursive(node.left, val)
                node.left.parent = node
            elif val > node.val:
                node.right = recursive(node.right, val)
                node.right.parent = node

            return node

        self.root = recursive(self.root, val)

    def inorder_traversal_recursive(self):
        def recursive(node):
            if node:
                recursive(node.left)
                result.append(node.val)
                recursive(node.right)

        result = []
        recursive(self.root)
        return result

    def preorder_traversal_recursive(self):
        def recursive(node):
            if node:
                result.append(node.val)
                recursive(node.left)
                recursive(node.right)

        result = []
        recursive(self.root)
        return result

    def postorder_traversal_recursive(self):
        def recursive(node):
            if node:
                recursive(node.left)
                recursive(node.right)
                result.append(node.val)

        result = []
        recursive(self.root)
        return result

思路

• 如果node有右子树，后继节点为右子树上最左边的节点
• 如果node没有右子树，假设node的父节点为p，如果node是p的左孩子，那么p就是node的后继节点；否则如果node是p的右孩子，则向上找，找到是子节点是父节点左孩子的节点。

实现

    def get_left_most_node(self, node):
        if node is None:
            return node

        while node.left is not None:
            node = node.left

        return node

    def successor(self, node):
        if node is None:
            return node

        if node.right is not None:
            return self.get_left_most_node(node.right)
        else:
            parent = node.parent
            while parent is not None and parent.left != node:
                node = parent
                parent = node.parent
            return parent

测试

import unittest
import random
import bst as bst
import operator


class FindSuccessorTestCase(unittest.TestCase):
    def test_find_successor_simple(self):
        count = 10
        data = [i for i in range(count)]
        random.shuffle(data)

        tree = bst.BST()
        for d in data:
            tree.insert_recursive(d)

        inorder = tree.inorder_traversal_recursive()
        node = tree.get_left_node(tree.root)
        successors = []
        while node:
            successors.append(node.val)
            node = tree.successor(node)

        self.assertTrue(operator.eq(inorder, successors))

    def test_find_successor(self):
        test_count = 100
        count = 100

        for _ in range(test_count):
            data = [i for i in range(count)]
            random.shuffle(data)
            tree = bst.BST()
            for d in data:
                tree.insert_recursive(d)

            inorder = tree.inorder_traversal_recursive()
            node = tree.get_left_node(tree.root)
            successors = []
            while node:
                successors.append(node.val)
                node = tree.successor(node)

            self.assertTrue(operator.eq(inorder, successors))

结果

➜  17_find_successor git:(master) ✗ python -m unittest test_find_successor.py
..
----------------------------------------------------------------------
Ran 2 tests in 0.038s

OK