Balance, and Sorted Collections in Python
30 min · Intermediate
Need help? Ask a tutor on WhatsApp (opens in a new tab)Everything so far costs O(h). Whether h is about log n or n depends only on the shape. This lesson shows how to get a good shape, and what Python offers when you want sorted keys without writing a tree.
Balance
Balanced tree is a tree whose shape stays close to the best possible: the heights of the two subtrees of every node differ by at most one. Its height is O(log n), so search, insert and delete are O(log n).
A tree with n keys has a height of at least log n. A balanced tree reaches that, so search, insert and delete cost O(log n). A chain has height n and costs O(n). Sorted input, the most common input, makes a chain.
Build a balanced tree from sorted keys
If all keys are known and sorted, take the middle key as the root. The left half builds the left subtree and the right half builds the right one. Each side holds about half the keys.
def balanced(keys, lo, hi): if lo > hi: return None mid = (lo + hi) // 2 node = Node(keys[mid]) node.left = balanced(keys, lo, mid - 1) node.right = balanced(keys, mid + 1, hi) return nodeExample 1 · A balanced tree from sorted keys
class Node: def __init__(self, key): self.key = key self.left = None self.right = None
def insert(node, key): if node is None: return Node(key) if key < node.key: node.left = insert(node.left, key) elif key > node.key: node.right = insert(node.right, key) return node
def height(node): if node is None: return 0 return 1 + max(height(node.left), height(node.right))
def balanced(keys, lo, hi): if lo > hi: return None mid = (lo + hi) // 2 node = Node(keys[mid]) node.left = balanced(keys, lo, mid - 1) node.right = balanced(keys, mid + 1, hi) return node
def inorder(node): if node is not None: yield from inorder(node.left) yield node.key yield from inorder(node.right)
keys = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]root = balanced(keys, 0, len(keys) - 1)print(root.key, height(root))print(*inorder(root))chain = Nonefor key in keys: chain = insert(chain, key)print(height(chain))5 4 1 2 3 4 5 6 7 8 9 10 10
What this code does
- The middle key 5 is the root, and the height is 4, which is the smallest possible for 10 keys.
- Inserting the same sorted keys one by one makes a chain of height 10. Same keys, very different cost.
Example 2 · Chain height against balanced height
class Node: def __init__(self, key): self.key = key self.left = None self.right = None
def insert(node, key): if node is None: return Node(key) if key < node.key: node.left = insert(node.left, key) elif key > node.key: node.right = insert(node.right, key) return node
def height(node): if node is None: return 0 return 1 + max(height(node.left), height(node.right))
def balanced(keys, lo, hi): if lo > hi: return None mid = (lo + hi) // 2 node = Node(keys[mid]) node.left = balanced(keys, lo, mid - 1) node.right = balanced(keys, mid + 1, hi) return node
import sys
sys.setrecursionlimit(5000)keys = list(range(1, 1001))chain = Nonefor key in keys: chain = insert(chain, key)good = balanced(keys, 0, len(keys) - 1)print(height(chain), height(good))1000 10
What this code does
- For 1000 keys the chain has height 1000, so a search may visit 1000 nodes. The balanced tree has height 10.
- Python stops deep recursion at about 1000 calls by default. A long chain is also a risk for crashes, not only for speed.
lo and hi are positions in the list, so no slices are copied. The time is O(n) and the height is about log n. Call it as balanced(keys, 0, len(keys) - 1).
Self-balancing trees in brief
When keys arrive one at a time, AVL and red-black trees fix the shape after each insert or delete with small local changes called rotations. Writing one is long, and interviews rarely ask for it.
Python has no built-in sorted map
Python's dict and set are hash tables with no key order. The standard library has no balanced tree. For a sorted list of keys, bisect finds a position in O(log n), and insort inserts at that position.
from bisect import bisect_left, bisect_right, insortkeys = [10, 20, 30]print(bisect_left(keys, 20))print(bisect_right(keys, 20))insort(keys, 25)print(keys)Example 3 · bisect for a sorted list
from bisect import bisect_left, bisect_right, insort
keys = []for key in [40, 10, 30, 20, 30]: insort(keys, key)print(keys)print(bisect_left(keys, 30), bisect_right(keys, 30))print(bisect_right(keys, 30) - bisect_left(keys, 30))[10, 20, 30, 30, 40] 2 4 2
What this code does
insortkeeps the list sorted, and equal keys are kept. The two positions around 30 give the count of 30s.- A set-like behaviour needs a check before inserting.
Example 4 · A small sorted set class
from bisect import bisect_left, bisect_right, insort
class SortedKeys: def __init__(self): self.keys = []
def add(self, key): i = bisect_left(self.keys, key) if i == len(self.keys) or self.keys[i] != key: self.keys.insert(i, key)
def floor(self, x): i = bisect_right(self.keys, x) return self.keys[i - 1] if i else None
def between(self, low, high): return self.keys[bisect_left(self.keys, low):bisect_right(self.keys, high)]
prices = SortedKeys()for price in [50, 20, 80, 20, 35]: prices.add(price)print(prices.keys)print(prices.floor(40), prices.floor(10))print(prices.between(30, 80))[20, 35, 50, 80] 35 None [35, 50, 80]
What this code does
- The class keeps unique keys sorted. It answers floor and range with two
bisectcalls. - This covers many problems that need a sorted map of keys, without writing a tree.
bisect_left gives the first position where the key could go, and bisect_right the one after equal keys. Search is O(log n). insort is O(n) because items shift, but the shift is a fast memory copy, so a list works well up to many thousands of keys.
The options and when to use each
A heap (heapq) gives only the smallest key fast. A sorted list with bisect suits mostly-static data. The third-party sortedcontainers package has SortedList and SortedDict for large changing data. A hand-written BST shows you understand the idea.
Next step
Check your understanding
11 quick questions