treap.py

from __future__ importannotations

fromrandomimportrandom


classNode:
"""
 Treap's node
 Treap is a binary tree by value and heap by priority
 """

def__init__(self, value: int|None=None):
self.value=value
self.prior=random()
self.left: Node|None=None
self.right: Node|None=None

def__repr__(self) ->str:
frompprintimportpformat

ifself.leftisNoneandself.rightisNone:
returnf"'{self.value}: {self.prior:.5}'"
else:
returnpformat(
 {f"{self.value}: {self.prior:.5}": (self.left, self.right)}, indent=1
 )

def__str__(self) ->str:
value=str(self.value) +" "
left=str(self.leftor"")
right=str(self.rightor"")
returnvalue+left+right


defsplit(root: Node|None, value: int) ->tuple[Node|None, Node|None]:
"""
 We split current tree into 2 trees with value:

 Left tree contains all values less than split value.
 Right tree contains all values greater or equal, than split value
 """
ifrootisNoneorroot.valueisNone: # None tree is split into 2 Nones
returnNone, None
elifvalue<root.value:
"""
 Right tree's root will be current node.
 Now we split(with the same value) current node's left son
 Left tree: left part of that split
 Right tree's left son: right part of that split
 """
left, root.left=split(root.left, value)
returnleft, root
else:
"""
 Just symmetric to previous case
 """
root.right, right=split(root.right, value)
returnroot, right


defmerge(left: Node|None, right: Node|None) ->Node|None:
"""
 We merge 2 trees into one.
 Note: all left tree's values must be less than all right tree's
 """
if (notleft) or (notright): # If one node is None, return the other
returnleftorright
elifleft.prior<right.prior:
"""
 Left will be root because it has more priority
 Now we need to merge left's right son and right tree
 """
left.right=merge(left.right, right)
returnleft
else:
"""
 Symmetric as well
 """
right.left=merge(left, right.left)
returnright


definsert(root: Node|None, value: int) ->Node|None:
"""
 Insert element

 Split current tree with a value into left, right,
 Insert new node into the middle
 Merge left, node, right into root
 """
node=Node(value)
left, right=split(root, value)
returnmerge(merge(left, node), right)


deferase(root: Node|None, value: int) ->Node|None:
"""
 Erase element

 Split all nodes with values less into left,
 Split all nodes with values greater into right.
 Merge left, right
 """
left, right=split(root, value-1)
_, right=split(right, value)
returnmerge(left, right)


definorder(root: Node|None) ->None:
"""
 Just recursive print of a tree
 """
ifnotroot: # None
return
else:
inorder(root.left)
print(root.value, end=",")
inorder(root.right)


definteract_treap(root: Node|None, args: str) ->Node|None:
"""
 Commands:
 + value to add value into treap
 - value to erase all nodes with value

 >>> root = interact_treap(None, "+1")
 >>> inorder(root)
 1,
 >>> root = interact_treap(root, "+3 +5 +17 +19 +2 +16 +4 +0")
 >>> inorder(root)
 0,1,2,3,4,5,16,17,19,
 >>> root = interact_treap(root, "+4 +4 +4")
 >>> inorder(root)
 0,1,2,3,4,4,4,4,5,16,17,19,
 >>> root = interact_treap(root, "-0")
 >>> inorder(root)
 1,2,3,4,4,4,4,5,16,17,19,
 >>> root = interact_treap(root, "-4")
 >>> inorder(root)
 1,2,3,5,16,17,19,
 >>> root = interact_treap(root, "=0")
 Unknown command
 """
forarginargs.split():
ifarg[0] =="+":
root=insert(root, int(arg[1:]))

elifarg[0] =="-":
root=erase(root, int(arg[1:]))

else:
print("Unknown command")

returnroot


defmain() ->None:
"""After each command, program prints treap"""
root=None
print(
"enter numbers to create a tree, + value to add value into treap, "
"- value to erase all nodes with value. 'q' to quit. "
 )

args=input()
whileargs!="q":
root=interact_treap(root, args)
print(root)
args=input()

print("good by!")


if__name__=="__main__":
importdoctest

doctest.testmod()
main()