From 179284a41bf265c59913b72d94bee47d7d6b7414 Mon Sep 17 00:00:00 2001 From: Hocnonsense <48747984+Hocnonsense@users.noreply.github.com> Date: Fri, 18 Oct 2019 15:39:37 +0800 Subject: [PATCH] Update treap.py (#1358) * Update treap.py check merge() * Update treap.py random() is used. its difficult to write doctests l->left r->right key->value add __repr__ and __str__ in preorder --- data_structures/binary_tree/treap.py | 186 +++++++++++++++++---------- 1 file changed, 115 insertions(+), 71 deletions(-) diff --git a/data_structures/binary_tree/treap.py b/data_structures/binary_tree/treap.py index 5d34abc3c..0b5947f4c 100644 --- a/data_structures/binary_tree/treap.py +++ b/data_structures/binary_tree/treap.py @@ -2,129 +2,173 @@ from random import random from typing import Tuple -class Node: +class Node(object): """ Treap's node - Treap is a binary tree by key and heap by priority + Treap is a binary tree by value and heap by priority """ - - def __init__(self, key: int): - self.key = key + def __init__(self, value: int = None): + self.value = value self.prior = random() - self.l = None - self.r = None + self.left = None + self.right = None + def __repr__(self): + from pprint import pformat -def split(root: Node, key: int) -> Tuple[Node, Node]: + if self.left is None and self.right is None: + return "'%s: %.5s'" % (self.value, self.prior) + else: + return pformat( + { + "%s: %.5s" + % (self.value, self.prior): (self.left, self.right) + }, + indent=1, + ) + + def __str__(self): + value = str(self.value) + " " + left = str(self.left or "") + right = str(self.right or "") + return value + left + right + +def split(root: Node, value: int) -> Tuple[Node, Node]: """ - We split current tree into 2 trees with key: + We split current tree into 2 trees with value: - Left tree contains all keys less than split key. - Right tree contains all keys greater or equal, than split key + Left tree contains all values less than split value. + Right tree contains all values greater or equal, than split value """ if root is None: # None tree is split into 2 Nones return (None, None) - if root.key >= key: - """ - Right tree's root will be current node. - Now we split(with the same key) current node's left son - Left tree: left part of that split - Right tree's left son: right part of that split - """ - l, root.l = split(root.l, key) - return (l, root) + elif root.value is None: + return (None, None) else: - """ - Just symmetric to previous case - """ - root.r, r = split(root.r, key) - return (root, r) - + if value < 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) + return (left, root) + else: + """ + Just symmetric to previous case + """ + root.right, right = split(root.right, value) + return (root, right) def merge(left: Node, right: Node) -> Node: """ We merge 2 trees into one. - Note: all left tree's keys must be less than all right tree's + Note: all left tree's values must be less than all right tree's """ - if (not left) or (not right): - """ - If one node is None, return the other - """ + if (not left) or (not right): # If one node is None, return the other return left or right - if left.key > right.key: + elif left.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.r = merge(left.r, right) + left.right = merge(left.right, right) return left else: """ Symmetric as well """ - right.l = merge(left, right.l) + right.left = merge(left, right.left) return right - -def insert(root: Node, key: int) -> Node: +def insert(root: Node, value: int) -> Node: """ Insert element - Split current tree with a key into l, r, + Split current tree with a value into left, right, Insert new node into the middle - Merge l, node, r into root + Merge left, node, right into root """ - node = Node(key) - l, r = split(root, key) - root = merge(l, node) - root = merge(root, r) - return root + node = Node(value) + left, right = split(root, value) + return merge(merge(left, node), right) - -def erase(root: Node, key: int) -> Node: +def erase(root: Node, value: int) -> Node: """ Erase element - Split all nodes with keys less into l, - Split all nodes with keys greater into r. - Merge l, r + Split all nodes with values less into left, + Split all nodes with values greater into right. + Merge left, right """ - l, r = split(root, key) - _, r = split(r, key + 1) - return merge(l, r) + left, right = split(root, value-1) + _, right = split(right, value) + return merge(left, right) - -def node_print(root: Node): +def inorder(root: Node): """ Just recursive print of a tree """ - if not root: + if not root: # None return - node_print(root.l) - print(root.key, end=" ") - node_print(root.r) + else: + inorder(root.left) + print(root.value, end=" ") + inorder(root.right) -def interactTreap(): +def interactTreap(root, args): """ Commands: - + key to add key into treap - - key to erase all nodes with key + + value to add value into treap + - value to erase all nodes with value - After each command, program prints treap + >>> root = interactTreap(None, "+1") + >>> inorder(root) + 1 + >>> root = interactTreap(root, "+3 +5 +17 +19 +2 +16 +4 +0") + >>> inorder(root) + 0 1 2 3 4 5 16 17 19 + >>> root = interactTreap(root, "+4 +4 +4") + >>> inorder(root) + 0 1 2 3 4 4 4 4 5 16 17 19 + >>> root = interactTreap(root, "-0") + >>> inorder(root) + 1 2 3 4 4 4 4 5 16 17 19 + >>> root = interactTreap(root, "-4") + >>> inorder(root) + 1 2 3 5 16 17 19 + >>> root = interactTreap(root, "=0") + Unknown command """ - root = None - while True: - cmd = input().split() - cmd[1] = int(cmd[1]) - if cmd[0] == "+": - root = insert(root, cmd[1]) - elif cmd[0] == "-": - root = erase(root, cmd[1]) + for arg in args.split(): + if arg[0] == "+": + root = insert(root, int(arg[1:])) + + elif arg[0] == "-": + root = erase(root, int(arg[1:])) + else: print("Unknown command") - node_print(root) + return root + +def main(): + """After each command, program prints treap""" + root = None + print("enter numbers to creat a tree, + value to add value into treap, - value to erase all nodes with value. 'q' to quit. ") + + args = input() + while args != 'q': + root = interactTreap(root, args) + print(root) + args = input() + + print("good by!") + pass if __name__ == "__main__": - interactTreap() + import doctest + doctest.testmod() + main()