bintree.py revision 2a99a7e74a7f215066514fe81d2bfa6639d9eddd 
1#!/usr/bin/env python
2#coding:utf-8
3# Author:  mozman
4# Purpose: binary tree module
5# Created: 28.04.2010
6# Copyright (c) 2010-2013 by Manfred Moitzi
7# License: MIT License
8
9from __future__ import absolute_import
10
11from .treemixin import TreeMixin
12
13__all__ = ['BinaryTree']
14
15
16class Node(object):
17    """ Internal object, represents a treenode """
18    __slots__ = ['key', 'value', 'left', 'right']
19
20    def __init__(self, key, value):
21        self.key = key
22        self.value = value
23        self.left = None
24        self.right = None
25
26    def __getitem__(self, key):
27        """ x.__getitem__(key) <==> x[key], where key is 0 (left) or 1 (right) """
28        return self.left if key == 0 else self.right
29
30    def __setitem__(self, key, value):
31        """ x.__setitem__(key, value) <==> x[key]=value, where key is 0 (left) or 1 (right) """
32        if key == 0:
33            self.left = value
34        else:
35            self.right = value
36
37    def free(self):
38        """ Set references to None """
39        self.left = None
40        self.right = None
41        self.value = None
42        self.key = None
43
44
45class BinaryTree(TreeMixin):
46    """
47    BinaryTree implements an unbalanced binary tree with a dict-like interface.
48
49    see: http://en.wikipedia.org/wiki/Binary_tree
50
51    A binary tree is a tree data structure in which each node has at most two
52    children.
53
54    BinaryTree() -> new empty tree.
55    BinaryTree(mapping,) -> new tree initialized from a mapping
56    BinaryTree(seq) -> new tree initialized from seq [(k1, v1), (k2, v2), ... (kn, vn)]
57
58    see also TreeMixin() class.
59
60    """
61    def __init__(self, items=None):
62        """ x.__init__(...) initializes x; see x.__class__.__doc__ for signature """
63        self._root = None
64        self._count = 0
65        if items is not None:
66            self.update(items)
67
68    def clear(self):
69        """ T.clear() -> None.  Remove all items from T. """
70        def _clear(node):
71            if node is not None:
72                _clear(node.left)
73                _clear(node.right)
74                node.free()
75        _clear(self._root)
76        self._count = 0
77        self._root = None
78
79    @property
80    def root(self):
81        """ root node of T """
82        return self._root
83
84    @property
85    def count(self):
86        """ count of items """
87        return self._count
88
89    def _new_node(self, key, value):
90        """ Create a new tree node. """
91        self._count += 1
92        return Node(key, value)
93
94    def insert(self, key, value):
95        """ T.insert(key, value) <==> T[key] = value, insert key, value into Tree """
96        if self._root is None:
97            self._root = self._new_node(key, value)
98        else:
99            parent = None
100            direction = 0
101            node = self._root
102            while True:
103                if node is None:
104                    parent[direction] = self._new_node(key, value)
105                    break
106                if key == node.key:
107                    node.value = value  # replace value
108                    break
109                else:
110                    parent = node
111                    direction = 0 if key <= node.key else 1
112                    node = node[direction]
113
114    def remove(self, key):
115        """ T.remove(key) <==> del T[key], remove item <key> from tree """
116        node = self._root
117        if node is None:
118            raise KeyError(str(key))
119        else:
120            parent = None
121            direction = 0
122            while True:
123                if key == node.key:
124                    # remove node
125                    if (node.left is not None) and (node.right is not None):
126                        # find replacment node: smallest key in right-subtree
127                        parent = node
128                        direction = 1
129                        replacement = node.right
130                        while replacement.left is not None:
131                            parent = replacement
132                            direction = 0
133                            replacement = replacement.left
134                        parent[direction] = replacement.right
135                        #swap places
136                        node.key = replacement.key
137                        node.value = replacement.value
138                        node = replacement  # delete replacement!
139                    else:
140                        down_dir = 1 if node.left is None else 0
141                        if parent is None:  # root
142                            self._root = node[down_dir]
143                        else:
144                            parent[direction] = node[down_dir]
145                    node.free()
146                    self._count -= 1
147                    break
148                else:
149                    direction = 0 if key < node.key else 1
150                    parent = node
151                    node = node[direction]
152                    if node is None:
153                        raise KeyError(str(key))
154