หาค่ามัธยฐานของ BST ใน O(n) time และ O(1) space ด้วย Python

โจทย์: หาค่ามัธยฐาน (median) ของ Binary Search Tree โดยใช้เวลา O(n) และพื้นที่ O(1) (Morris Inorder Traversal)

นิยาม: จำนวนโหนดคี่ → median = โหนดที่ (n+1)/2, จำนวนโหนดคู่ → median = (โหนด n/2 + โหนด n/2+1) / 2

อัลกอริทึม Morris Inorder Traversal

1. นับจำนวนโหนดทั้งหมด (node_count) ด้วย Morris Traversal
2. ทำ Morris Traversal อีกครั้ง นับลำดับ (count_curr)
3. เมื่อถึงตำแหน่ง median → คืนค่า (หรือค่าเฉลี่ย 2 ตัวถ้าคู่)

Morris Traversal: ใช้ threaded binary tree ไม่ต้องใช้ stack/recursion → O(1) space

ตัวอย่างโค้ด

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

def number_of_nodes(root):
    count = 0
    current = root
    while current:
        if current.left is None:
            count += 1
            current = current.right
        else:
            pre = current.left
            while pre.right and pre.right != current:
                pre = pre.right
            if pre.right is None:
                pre.right = current
                current = current.left
            else:
                pre.right = None
                count += 1
                current = current.right
    return count

def calculate_median(root):
    if not root: return 0
    node_count = number_of_nodes(root)
    count_curr = 0
    current = root
    prev = None
    while current:
        if current.left is None:
            count_curr += 1
            if node_count % 2 == 1 and count_curr == (node_count + 1)//2:
                return prev.data if prev else current.data
            elif node_count % 2 == 0 and count_curr == (node_count//2) + 1:
                return (prev.data + current.data) // 2
            prev = current
            current = current.right
        else:
            pre = current.left
            while pre.right and pre.right != current:
                pre = pre.right
            if pre.right is None:
                pre.right = current
                current = current.left
            else:
                pre.right = None
                count_curr += 1
                if node_count % 2 == 1 and count_curr == (node_count + 1)//2:
                    return current.data
                elif node_count % 2 == 0 and count_curr == (node_count//2) + 1:
                    return (prev.data + current.data) // 2
                prev = current
                current = current.right

# ตัวอย่างต้นไม้
root = TreeNode(7)
root.left = TreeNode(4)
root.right = TreeNode(9)
root.left.left = TreeNode(2)
root.left.right = TreeNode(5)
root.right.left = TreeNode(8)
root.right.right = TreeNode(10)

print(calculate_median(root))  # 7

ผลลัพธ์

7

สรุป: Morris Traversal ช่วยให้ inorder traversal ไม่ต้องใช้ stack ได้ O(1) space เหมาะสำหรับหา median, kth smallest ฯลฯ