treap.py

import random


class Node:
    def __init__(self, key, priority):
        self.key = key
        self.priority = priority
        self.parent = None
        self.left = None
        self.right = None


class Treap:
    def __init__(self):
        self.root = None

    def __iter__(self):
        def visit(node):
            if node.left is not None:
                yield from visit(node.left)
            yield node.key
            if node.right is not None:
                yield from visit(node.right)

        if self.root is not None:
            yield from visit(self.root)

    def rotate_left(self, x):
        y = x.right  # yをxの右の子とする
        x.right = y.left  # yの左部分木をxの右部分木にする
        if y.left is not None:
            y.left.parent = x
        y.parent = x.parent  # xの親をyの親にする
        if x.parent is None:  # xが根の時
            self.root = y
        elif x == x.parent.left:  # xが親の左部分木の時
            x.parent.left = y
        else:  # xが親の右部分木の時
            x.parent.right = y
        y.left = x  # xをyの右の子にする
        x.parent = y  # xの親をyにする

    def rotate_right(self, y):
        x = y.left
        y.left = x.right
        if x.right is not None:
            x.right.parent = y
        x.parent = y.parent
        if y.parent is None:
            self.root = x
        elif y == y.parent.right:
            y.parent.right = x
        else:
            y.parent.left = x
        x.right = y
        y.parent = x

    def insert(self, key):
        z = Node(key, random.random())
        y = None
        # 根から挿入先を探索する
        x = self.root
        while x is not None:
            y = x
            if z.key < x.key:  # z,keyがx.keyより小さければxの左部分木を探索
                x = x.left
            else:  # z.keyがx.keyより大きければxの右部分木を探索
                x = x.right
        z.parent = y
        if y is None:
            self.root = z
        elif z.key < y.key:
            y.left = z
        else:
            y.right = z
        self.trickle_up(z)

    def trickle_up(self, z):
        while z.parent is not None and z.parent.priority > z.priority:
            if z.parent.right == z:
                self.rotate_left(z.parent)
            else:
                self.rotate_right(z.parent)
        if z.parent is None:
            self.root = z

    def delete(self, key):
        z = self.find_node(key)
        if z is not None and z.key == key:
            self.trickle_down(z)
            if z.parent.right == z:
                z.parent.right = None
            else:
                z.parent.left = None

    def trickle_down(self, z):
        while z.left is not None or z.right is not None:
            if z.left is None:
                self.rotate_left(z)
            elif z.right is None:
                self.rotate_right(z)
            elif z.left.priority < z.right.priority:
                self.rotate_right(z)
            else:
                self.rotate_left(z)
        if z == self.root:
            self.root = z.parent

    def find_node(self, key):
        node = self.root
        while node is not None:
            if node.key == key:
                return node
            if key < node.key:
                node = node.left
            else:
                node = node.right
        return node