2分探索木

「プログラミングの基礎」をちょっとずつ読んでて2分探索木の走査がでてきたんで、pythonでも書いてみました。

ocamlだとこう

type tree_t = Empty
	      | Leaf of int
	      | Node of tree_t * int * tree_t

(* 目的 : dataが2分探索木treeに含まれているかを調べる *)
(* search : tree_t -> int -> bool *)
let rec search tree  data = match tree with
    Empty -> false
  | Leaf (n) -> data = n
  | Node (t1, n, t2) ->
    if data = n
    then true
    else
      if data < n
      then search t1 data
      else search t2 data

目的のdataがNodeの値よりも小さければ木構造の左側を、大きければ右側を返すというのを再帰で繰り返してます。

そのままpythonで書くとこんな感じかな

class Node:
    def __init__(self, n=None, left=None, right=None):
        self.data = n
        self.left = left
        self.right = right

    @staticmethod
    def make_empty():
        return Node()

    @staticmethod
    def make_leaf(n):
        return Node(n)


def search(tree, data):
    """dataが2分探索木treeに含まれているかを調べる"""
    # Empty
    if tree.data is None and tree.left is None and tree.right is None:
        return False
    # Leaf
    elif tree.data is not None and tree.left is None and tree.right is None:
        return tree.data == data
    # Node
    else:
        if tree.data == data:
            return True
        elif data < tree.data:
            return search(tree.left, data)
        else:
            return search(tree.right, data)

# data
tree1 = Node.make_empty()
tree2 = Node.make_leaf(3)
tree3 = Node(left=Node.make_leaf(1), n=2, right=Node.make_leaf(3))
tree4 = Node(left=Node.make_empty(), n=7, right=Node.make_leaf(9))
tree5 = Node(left=tree3, n=6, right=tree4)


# test
assert not search(tree1, 3)
assert search(tree2, 3)
assert not search(tree2, 4)
assert search(tree5, 6)
assert search(tree5, 2)
assert search(tree5, 1)
assert not search(tree5, 4)

再帰が深くなるとエラーになるので再帰を使わないバージョン

def search(tree, data):
    """dataが2分探索木treeに含まれているかを調べる"""
    target_tree = tree
    while True:
        # Empty
        if target_tree.data is None and target_tree.left is None and target_tree.right is None:
            return False
        # Leaf
        elif target_tree.data is not None and target_tree.left is None and target_tree.right is None:
            return target_tree.data == data
        # Node
        else:
            if target_tree.data == data:
                return True
            elif data < target_tree.data:
                target_tree = target_tree.left
            else:
                target_tree = target_tree.right

プログラミングの基礎 (Computer Science Library)