Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
279 changes: 279 additions & 0 deletions data_structures/binary_tree/b_tree.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,279 @@
"""
B-Tree Implementation

A B-Tree is a self-balancing tree data structure that maintains sorted data and allows
searches, sequential access, insertions, and deletions in logarithmic time.

B-Trees are commonly used in databases and file systems.

Reference: https://en.wikipedia.org/wiki/B-tree
Time Complexity:
- Search: O(log n)
- Insert: O(log n)
- Delete: O(log n)
"""

from __future__ import annotations


class BTreeNode:
"""
A node in the B-Tree.

Attributes:
keys: List of keys stored in the node
children: List of child nodes
is_leaf: Boolean indicating if this is a leaf node
"""

def __init__(self, is_leaf: bool = True) -> None:
self.keys: list[int] = []
self.children: list[BTreeNode] = []
self.is_leaf = is_leaf

def split(self, parent: BTreeNode, index: int) -> None:
"""
Split this node and move the median key up to the parent.

Args:
parent: The parent node
index: The index in parent's children where this node is located
"""
new_node = BTreeNode(is_leaf=self.is_leaf)
mid_index = len(self.keys) // 2
median_key = self.keys[mid_index]

new_node.keys = self.keys[mid_index + 1 :]
self.keys = self.keys[:mid_index]

if not self.is_leaf:
new_node.children = self.children[mid_index + 1 :]
self.children = self.children[: mid_index + 1]

parent.keys.insert(index, median_key)
parent.children.insert(index + 1, new_node)


class BTree:
"""
B-Tree data structure.

A B-Tree of order m has the following properties:
- Every node has at most m children
- Every non-leaf node (except root) has at least ⌈m/2⌉ children
- The root has at least 2 children if it is not a leaf
- All leaves appear on the same level
- A non-leaf node with k children contains k−1 keys

Check failure on line 66 in data_structures/binary_tree/b_tree.py

View workflow job for this annotation

GitHub Actions / ruff

ruff (RUF002)

data_structures/binary_tree/b_tree.py:66:49: RUF002 Docstring contains ambiguous `−` (MINUS SIGN). Did you mean `-` (HYPHEN-MINUS)?

Examples:
>>> btree = BTree(order=3)
>>> btree.insert(10)
>>> btree.insert(20)
>>> btree.insert(5)
>>> btree.insert(6)
>>> btree.insert(12)
>>> btree.insert(30)
>>> btree.insert(7)
>>> btree.insert(17)
>>> btree.search(6)
True
>>> btree.search(15)
False
>>> btree.search(12)
True
>>> btree.search(100)
False
"""

def __init__(self, order: int = 3) -> None:
"""
Initialize a B-Tree.

Args:
order: The maximum number of children a node can have (must be >= 3)

Raises:
ValueError: If order is less than 3
"""
if order < 3:
msg = "Order must be at least 3"
raise ValueError(msg)

self.order = order
self.min_keys = (order + 1) // 2 - 1
self.max_keys = order - 1
self.root = BTreeNode()

def search(self, key: int, node: BTreeNode | None = None) -> bool:
"""
Search for a key in the B-Tree.

Args:
key: The key to search for
node: The node to start searching from (defaults to root)

Returns:
True if the key exists, False otherwise

Time Complexity: O(log n)

>>> btree = BTree(order=3)
>>> btree.insert(50)
>>> btree.search(50)
True
>>> btree.search(25)
False
"""
if node is None:
node = self.root

i = 0
while i < len(node.keys) and key > node.keys[i]:
i += 1

if i < len(node.keys) and key == node.keys[i]:
return True

if node.is_leaf:
return False

return self.search(key, node.children[i])

def insert(self, key: int) -> None:
"""
Insert a key into the B-Tree.

Args:
key: The key to insert

Time Complexity: O(log n)

>>> btree = BTree(order=3)
>>> btree.insert(10)
>>> btree.insert(20)
>>> btree.insert(30)
>>> btree.search(20)
True
"""
if len(self.root.keys) >= self.max_keys:
new_root = BTreeNode(is_leaf=False)
new_root.children.append(self.root)
self.root.split(new_root, 0)
self.root = new_root

self._insert_non_full(self.root, key)

def _insert_non_full(self, node: BTreeNode, key: int) -> None:
"""
Insert a key into a node that is not full.

Args:
node: The node to insert into
key: The key to insert
"""
i = len(node.keys) - 1

if node.is_leaf:
node.keys.append(0)
while i >= 0 and key < node.keys[i]:
node.keys[i + 1] = node.keys[i]
i -= 1
node.keys[i + 1] = key
else:
while i >= 0 and key < node.keys[i]:
i -= 1
i += 1

if len(node.children[i].keys) >= self.max_keys:
node.children[i].split(node, i)
if key > node.keys[i]:
i += 1

self._insert_non_full(node.children[i], key)

def traverse(self, node: BTreeNode | None = None) -> list[int]:
"""
Traverse the B-Tree in sorted order.

Args:
node: The node to start traversal from (defaults to root)

Returns:
List of all keys in sorted order

>>> btree = BTree(order=3)
>>> for i in [10, 20, 5, 6, 12, 30, 7, 17]:
... btree.insert(i)
>>> btree.traverse()
[5, 6, 7, 10, 12, 17, 20, 30]
"""
if node is None:
node = self.root

result: list[int] = []

for i in range(len(node.keys)):
if not node.is_leaf and i < len(node.children):
result.extend(self.traverse(node.children[i]))
result.append(node.keys[i])

if not node.is_leaf and len(node.children) > len(node.keys):
result.extend(self.traverse(node.children[len(node.keys)]))

return result

def get_height(self, node: BTreeNode | None = None) -> int:
"""
Get the height of the B-Tree.

Args:
node: The node to start from (defaults to root)

Returns:
The height of the tree

>>> btree = BTree(order=3)
>>> btree.get_height()
0
>>> btree.insert(10)
>>> btree.get_height()
0
>>> for i in range(20):
... btree.insert(i)
>>> btree.get_height() > 0
True
"""
if node is None:
node = self.root

if node.is_leaf:
return 0

return 1 + self.get_height(node.children[0])

def __str__(self) -> str:
"""
String representation of the B-Tree.

Returns:
String showing all keys in sorted order
"""
return f"BTree(order={self.order}, keys={self.traverse()})"


if __name__ == "__main__":
import doctest

doctest.testmod()

btree = BTree(order=3)
keys = [10, 20, 5, 6, 12, 30, 7, 17, 3, 8, 15, 25, 35, 40]

print("Inserting keys:", keys)
for key in keys:
btree.insert(key)

print("\nB-Tree traversal (sorted):", btree.traverse())
print("B-Tree height:", btree.get_height())
print("\nSearching for 12:", btree.search(12))
print("Searching for 100:", btree.search(100))
Loading
Loading