github.com/benz9527/toy-box/algo@v0.0.0-20240221120937-66c0c6bd5abd/tree/treap.go (about)

     1  package tree
     2  
     3  // Binary Sort Tree + Heap = Treap（树堆）
     4  // Treap 不等同于 Binary Heap
     5  // 1. Binary Heap 必须是一个完全二叉树， Treap 则不一定是完全二叉树
     6  // 2. Treap 不严格满足 AVL Tree 的要求，因为 Treap 左右子树的高度差绝对值可能会超过 1
     7  //  只能近似满足 AVL Tree 的性质。下面的实现也能看出，进行调整的旋转只进行了一次，而没有
     8  //  AVL Tree 的连续调整。Treap 的旋转调整是根据优先级完成的，就是为了实现随机平衡结构，
     9  //  也能实现快速查找。
    10  
    11  type TreapNode struct {
    12  	Left, Right *TreapNode
    13  	// @Field Priority 随机优先级
    14  	Priority, Key int
    15  }
    16  
    17  func (t *TreapNode) RotateLeft() *TreapNode {
    18  	// FIXME(Ben) nil ptr
    19  	r := t.Right
    20  	t.Right = r.Left
    21  	r.Left = t
    22  	return r
    23  }
    24  
    25  func (t *TreapNode) RotateRight() *TreapNode {
    26  	// FIXME(Ben) nil ptr
    27  	l := t.Left
    28  	t.Left = l.Right
    29  	l.Right = t
    30  	return l
    31  }
    32  
    33  type Treap struct {
    34  	Root *TreapNode
    35  }
    36  
    37  func (t *Treap) Insert(key, priority int) {
    38  	t.insert(t.Root, key, priority)
    39  }
    40  
    41  func (t *Treap) insert(root *TreapNode, key, priority int) {
    42  	if root == nil {
    43  		root = &TreapNode{
    44  			Left:     nil,
    45  			Right:    nil,
    46  			Key:      key,
    47  			Priority: priority,
    48  		}
    49  		return
    50  	}
    51  	// classical dichotomy & recursive
    52  	if key < root.Key {
    53  		// Insert into left node
    54  		t.insert(root.Left, key, priority)
    55  		// rotate adjust
    56  		if root.Left.Priority < root.Priority {
    57  			root = root.RotateRight()
    58  		}
    59  	} else {
    60  		// Insert into right node
    61  		t.insert(root.Right, key, priority)
    62  		// rotate adjust
    63  		if root.Right.Priority < root.Priority {
    64  			root = root.RotateLeft()
    65  		}
    66  	}
    67  }
    68  
    69  func (t *Treap) Delete(key int) {
    70  	t.delete(t.Root, key)
    71  }
    72  
    73  func (t *Treap) delete(root *TreapNode, key int) {
    74  	if root == nil {
    75  		return
    76  	}
    77  
    78  	if key < root.Key {
    79  		t.delete(root.Left, key)
    80  	} else if key > root.Key {
    81  		t.delete(root.Right, key)
    82  	} else {
    83  		if root.Left == nil {
    84  			root = root.Right
    85  		} else if root.Right == nil {
    86  			root = root.Left
    87  		} else {
    88  			if root.Left.Priority < root.Right.Priority {
    89  				root = root.RotateRight()
    90  				t.delete(root.Right, key)
    91  			} else {
    92  				root = root.RotateLeft()
    93  				t.delete(root.Left, key)
    94  			}
    95  		}
    96  	}
    97  }