AVL Tree ใน python

นั่งแอบเขียน AVL Tree ระหว่างนั่งเรียนวิชา Abstract Data Types (ตัวอย่างไม่ดี โปรดอย่าเลียนแบบ :P)

AVL Tree คือ binary search tree ที่คอยจัดตัวเองให้ balance อยู่เสมอ (subtree ทางซ้าย กับ subtree ทางขวามี depth แตกต่างกันไม่เกิน 2) และการจัดต้นไม้ให้ balance อยู่เสมอทำให้ binary search tree สามารถแทรก/ค้นหา ข้อมูลได้อย่างมีประสิทธิภาพ (ไม่มีปัญหาที่เกิดจากต้นไม้เอียงไปข้างใดข้างหนึ่ง)

AVL Tree ที่เขียนเป็น class มี method isLeaf (เช็คว่าเป็นโหนดลูก), depth (หาความลึกของต้นไม้), balance (คำนวณ balance factor สำหรับต้นไม้), LeftRotate (ทำการหมุนต้นไม้ไปทางซ้าย), RightRotate (หมุนต้นไม้ไปทางขวา), RLDoubleRotate (หมุนสองครั้งแบบขวา-ซ้าย), LRDoubleRotate (หมุนสองครั้งแบบซ้าย-ขวา), insert (แทรกโหนดลงไปในกราฟ), travel (ท่องไปในต้นไม้), printtree (พิมพ์ต้นไม้)

ตัวอย่างต้นไม้ที่ printtree พิมพ์ออกมา

    7
  6
    5
4
    3
  2
    1

ข้างบนเป็นต้นไม้ที่มี 4 เป็น root และมี 2 กับ 6 เป็นลูกทางซ้ายและขวาตามลำดับ โดย 2 มี 1 และ 3 เป็นลูก…. วิธี printtree แบบนี้ ได้เทคนิคมาจาก อ.ที่ค่าย สอวน. ศูนย์ ม.ศิลปากร ครับ 😀

(โค้ดอยู่ด้านในนะครับ)

# author: Natt Piyapramote 
class AVLTree:
    def __init__(self, data):
        self.data = data
        self.left = None
        self.right = None

    def isLeaf(self):
        if self.left is None and self.right is None:
            return True
        return False

    def depth(self):
        if self.isLeaf():
            return 0
        depth = 0
        if self.left is not None:
            depth = self.left.depth()
        if self.right is not None:
            depth = max(self.right.depth(), depth)
        return depth+1

    def balance(self):
        if self.isLeaf():
            return 0
        if self.right is None:
            return self.left.depth()+1
        if self.left is None:
            return -(self.right.depth()+1)
        return self.left.depth() - self.right.depth()

    def LeftRotate(self, root):
        pivot = root.right
        root.right = pivot.left
        pivot.left = root
        return pivot

    def RightRotate(self, root):
        pivot = root.left
        root.left = pivot.right
        pivot.right = root
        return pivot

    def RLDoubleRotate(self, root):
        root.right = self.RightRotate(root.right)
        return self.LeftRotate(root)

    def LRDoubleRotate(self, root):
        root.left = self.LeftRotate(root.left)
        return self.RightRotate(root)

    def insert(self, node):
        # insert node
        if node.data  0:
                return self.RightRotate(self)
            else:
                return self.LRDoubleRotate(self)
        elif self.balance() == -2:
            if self.right and self.right.balance() < 0:
                return self.LeftRotate(self)
            else:
                return self.RLDoubleRotate(self)
        return self

    def travel(self):
        if self.left is not None:
            self.left.travel()
        print self.data, 
        if self.right is not None:
            self.right.travel()

    def printtree(self, depth=0, padding=2):
        "pre order travel to print binary tree"
        if self.right: self.right.printtree(depth+1)
        print ' '*padding*depth + `self.data`
        if self.left: self.left.printtree(depth+1)

root = AVLTree(4)
for i in [6, 1, 2, 3, 7, 5]:
    root = root.insert(AVLTree(i))
root.printtree()