class Node:
"""Conventions: a nonexistent child or parent is None."""
parent = None
leftchild = None
rightchild = None
annotation = None
value = None
class AANode(Node):
level = 1
class Walker:
"""descend must return 0 if the node in question is the one desired,
-1 if the left child would be better, or 1 if the right child would be
better"""
def descend(self, node):
raise
def prepare_descend(self, *args): #optional method to prepare for descent
raise
"""ascend should return True iff it should run again on the parent.
It should set the state such that a subsequent
descent would retrace these steps."""
def ascend(self, node):
raise
def prepare_ascend(self, *args): #optional method to prepare for ascent
raise
class SearchWalker(Walker):
"""Convention: leftchild.annotation < annotation < rightchild.annotation"""
val = 0
compare = cmp
def prepare_descend(self, val, comparator=cmp):
self.val = val
self.compare = comparator
def descend(self, node):
x = self.compare(node.annotation, self.val)
return x
def ascend(self, node):
self.val = node.annotation
return False
class RandomWalker(Walker):
def prepare_descend(self):
from random import choice as choice
self.choice = choice
def descend(self, node):
return self.choice((-1,1))
#ascend not implemented; it doesn't make sense because there is no
#state and no reproducibility
def descend(node, walker): #move down from a root node
x = walker.descend(node)
while x != 0:
if x == 1:
if node.rightchild is None:
return (node, 1)
else:
node = node.rightchild
else: #x == -1
if node.leftchild is None:
return (node, -1)
else:
node = node.leftchild
x = walker.descend(node)
return (node, 0)
def ascend(node, walker): #Move up from a leaf node
while node is not None and walker.ascend(node):
node = node.parent
def search(root, val):
"""Searches a correctly sorted binary tree, starting with the root Node, for
val. Returns a node that contains val if val is present, otherwise a node
for which val would be an acceptable child value."""
w = SearchWalker
w.prepared_descend(val)
return descend(root, w)
def findmin(root):
while root.leftchild is not None:
root = root.leftchild
return root
def findmax(root):
while root.rightchild is not None:
root = root.rightchild
return root
class MonoidTree:
makenode = Node
"""A Monoid Annotation Tree is a binary tree whose nodes are each annotated
by values from some monoid. The annotation of an internal node is computed
by applying the operation to the annotations of its children. The annotation of a leaf
node is specified by the user. Every node must either have two children or
be a leaf node.
Each leaf node may also be associated with an arbitrary opaque value of the user's
choosing. This node and value will remain associated."""
def __init__(self, operation, rootnode):
"""The rootnode must have a valid annotation, and its parent must be None"""
self.op = operation
self.root = rootnode
def _update(self, node, sentinel=None):
"""node must be an internal node"""
while node is not sentinel:
#oldval = node.annotation
node.annotation = self.op(node.leftchild.annotation, node.rightchild.annotation)
#if oldval == node.annotation:
# #this node has not changed, so nodes above it will also not have changed
# break
#else:
node = node.parent
_update_add = _update
_update_del = _update
def _split_link(self, node):
"""Introduce and return a new node (newparent) between node and its parent"""
newparent = self.makenode()
newparent.parent = node.parent
if node.parent is not None:
if node.parent.leftchild is node:
node.parent.leftchild = newparent
else:
assert node.parent.rightchild is node
node.parent.rightchild = newparent
else:
self.root = newparent
node.parent = newparent
return newparent
def addleft(self, new, old):
"""Add a new leaf node to the left of an old leaf node"""
newparent = self._split_link(old)
newparent.rightchild = old
newparent.leftchild = new
new.parent = newparent
self._update_add(newparent)
def addright(self, new, old):
"""Add a new leaf node to the right of an old leaf node"""
newparent = self._split_link(old)
newparent.rightchild = new
newparent.leftchild = old
new.parent = newparent
self._update_add(newparent)
def add(self, new, walker):
leaf, position = descend(self.root, walker)
assert leaf.leftchild is None
assert leaf.rightchild is None
if position == 1:
self.addright(new, leaf)
else: #Makes left the default for duplicate values
self.addleft(new, leaf)
def remove(self, leaf):
p = leaf.parent
if p.leftchild is leaf:
sibling = p.rightchild
else:
assert p.rightchild is leaf
sibling = p.leftchild
gp = p.parent
if gp.leftchild is p:
gp.leftchild = sibling
elif gp.rightchild is p:
gp.rightchild = sibling
sibling.parent = gp
# The only remaining reference to p is now in leaf itself, and the only
# remaining reference to leaf is in the user's hands
self._update_del(gp)
def change_annotation(self, leaf, newann):
assert leaf.leftchild is None
assert leaf.rightchild is None
leaf.annotation = newann
self._update(leaf.parent)
def getnext(self, leaf, skip=None):
assert leaf.leftchild is None
assert leaf.rightchild is None
node = leaf
while ((node.parent is not None) and
((node.parent.rightchild is node) or
((skip is not None) and skip(node.parent.rightchild)))):
# Move up until you can move right
node = node.parent
if (node.parent is not None) and (node.parent.leftchild is node):
node = node.parent.rightchild
while node.leftchild is not None:
# Move down, staying as far left as possible.
assert node.rightchild is not None
if (skip is not None) and skip(node.leftchild):
node = node.rightchild
else:
node = node.leftchild
return node
else:
raise StopIteration("No next node")
def _build_subtree(self, nodes):
#FIXME: This cannot be helpful because insertion of a subtree requires
#rebalancing the main tree by more than one level, which is not possible
#with a single invocation of skew and split
L = len(nodes)
if L == 1:
return nodes[0]
else:
next = []
sentinel = 'g' #must not be None, since None is the root sentinel
if L % 2:
n2 = nodes.pop()
n1 = nodes.pop()
newnode = self.makenode()
newnode.parent=sentinel #totally arbitrary constant
newnode.leftchild = n1
n1.parent = newnode
newnode.rightchild = n2
n2.parent = newnode
self._update_add(newnode, sentinel)
nodes.append(newnode)
for i in xrange(0,L,2):
n1,n2 = nodes[i:(i+2)]
newnode = self.makenode()
newnode.parent=sentinel #totally arbitrary constant
newnode.leftchild = n1
n1.parent = newnode
newnode.rightchild = n2
n2.parent = newnode
self._update_add(newnode, sentinel)
class SumWalker(Walker):
"""SumWalker is designed to walk over full trees where each leaf has annotation 1
and the monoid is +. Target is the zero-indexed position of the target node.
There is one exception: the last node in every tree has annotation 0."""
target = None
offset = None
def prepare_descend(self, target):
self.target = target
self.offset = 0
def descend(self, node):
if node.annotation == 0: #empty leaf at the last position
assert self.target == self.offset
return -1
elif node.leftchild is None: #leaf node case
assert node.rightchild is None
assert self.target == self.offset
return 0
else: #internal node case
p = self.offset + node.leftchild.annotation
if p <= self.target:
self.offset = p
return 1
else:
return -1
def prepare_ascend(self):
self.target = 0
def ascend(self, node):
if node.parent is not None:
if node.parent.rightchild is node:
self.target += node.parent.leftchild.annotation
else:
assert node.parent.leftchild is node
return True
else:
return False
class TreeList:
"""Implements a list-like interface, backed by a MonoidTree"""
_treetype = MonoidTree
def __init__(self):
self._makenode = self._treetype.makenode
r = self._makenode()
r.annotation = 0
from operator import add
self._tree = self._treetype(add, r)
self._walker = SumWalker()
# We regard the fields of this walker as public API, and manipulate
# them directly
self._index = {}
def __len__(self):
return self._tree.root.annotation
def _getnode(self, i):
self._walker.prepare_descend(i)
node, pos = descend(self._tree.root, self._walker)
assert pos == 0
return node
def __getitem__(self, s):
if isinstance(s, int):
node = self._getnode(s)
return node.value
else:
raise UnimplementedError
def __setitem__(self, s, v):
if isinstance(s, int):
if s < len(self):
node = self._getnode(s)
oldv = node.value
self._index[oldv].remove(node)
if not self._index[oldv]:
del self._index[oldv]
node.value = v
if v not in self._index:
self._index[v] = set()
self._index[v].add(node)
else:
self.insert(s, v)
else:
raise UnimplementedError
def __delitem__(self, s):
if isinstance(s, int):
if s < len(self):
node = self._getnode(s)
oldv = node.value
self._index[oldv].remove(node)
if not self._index[oldv]:
del self._index[oldv]
self._tree.remove(node)
else:
raise UnimplementedError
def insert(self, p, v):
if p > len(self):
raise IndexError("Index out of range")
self._walker.prepare_descend(p)
newnode = self._makenode()
newnode.annotation = 1
newnode.value = v
self._tree.add(newnode, self._walker)
if v not in self._index:
self._index[v] = set()
self._index[v].add(newnode)
def index(self, v):
"""index returns some index such that self[i] == v. No promises about ordering."""
self._walker.prepare_ascend()
for node in self._index[v]: #Pull one arbitrary node out of the set
assert node.value == v
ascend(node, self._walker)
break
return self._walker.target
class TreeHideList:
"""Implements the EagerHideList interface, backed by a MonoidTree"""
_treetype = MonoidTree
class MultiSumWalker(Walker):
index = 0
target = 0
offset = 0
def prepare_descend(self, target, index):
self.index = index
self.target = target
self.offset = 0
def descend(self, node):
if node.annotation == (0,0): #empty leaf at the last position
assert self.target == self.offset
return -1
elif node.leftchild is None: #leaf node case
assert node.rightchild is None
assert self.target == self.offset
return 0
else: #internal node case
p = self.offset + node.leftchild.annotation[self.index]
if p <= self.target:
self.offset = p
return 1
else:
return -1
def prepare_ascend(self, index):
self.target = 0
self.index = index
def ascend(self, node):
if node.parent is not None:
if node.parent.rightchild is node:
self.target += node.parent.leftchild.annotation[self.index]
else:
assert node.parent.leftchild is node
return True
else:
return False
@staticmethod
def op(a,b):
# Convention: a[0] is visible elements. a[1] is all elements.
return (a[0] + b[0], a[1] + b[1])
@staticmethod
def skip(node):
return node.annotation[0] == 0
def __init__(self):
self._makenode = self._treetype.makenode
r = self._makenode()
r.annotation = (0, 0)
self._tree = self._treetype(self.op, r)
self._walker = self.MultiSumWalker()
# We regard the fields of this walker as public API, and manipulate
# them directly
self._index = {}
unique = True
if unique:
self._index_lookup = self._index.__getitem__
self._index_assign = self._index.__setitem__
else:
self._index_lookup = self._index_lookup_set
self._index_assign = self._index_assign_set
def _index_lookup_set(self, item):
for v in self._index[item]:
return v
def _index_assign_set(self, key, value):
if key not in self._index:
self._index[key] = set()
self._index[key].add(value)
def __len__(self):
return self._tree.root.annotation[0]
def _getnode(self, i, a):
self._walker.prepare_descend(i, a)
node, pos = descend(self._tree.root, self._walker)
assert (pos == 0) or ((pos == -1) and (i == len(self)))
return node
def __getitem__(self, s):
if isinstance(s, int):
if s < len(self): #FIXME: negative indices
node = self._getnode(s, 0)
return node.value
else:
raise IndexError("Index out of range")
else:
start, stop, stride = s.indices(len(self))
if start == stop:
return []
elif stride == 1:
# runs in k + log(N) (amortized)
nodes = [self._getnode(start,0)]
k = stop - start
while len(nodes) < k:
nodes.append(self._tree.getnext(nodes[-1],self.skip))
return [n.value for n in nodes]
else:
#FIXME: runs in k*log(N), could be reduced to k*log(step) + log(N)
return [self[i] for i in xrange(start,stop,stride)]
def index(self, v, visible=True):
"""index returns some index such that self[i] == v. No promises about ordering."""
self._walker.prepare_ascend(0 if visible else 1)
node = self._index_lookup(v) #Pull one arbitrary node out of the set
assert node.value == v
ascend(node, self._walker)
return self._walker.target
def hide(self, position, length):
#self.__getitem__ is eager, so we acquire the list of nodes before
#acting on them
node = self._getnode(position,0)
for i in xrange(position+1,position+length):
self._tree.change_annotation(node,(0,1))
node = self._tree.getnext(node, self.skip)
self._tree.change_annotation(node,(0,1))
#FIXME: runs in length*log(N). Could be reduced using a priority queue,
#possibly to length + log(N)
def getitem_all(self, s):
if isinstance(s, int):
node = self._getnode(s, 1)
return node.value
else:
#FIXME: runs in k*log(N), could be reduced to k + log(N) by linked list
return [self.getitem_all(i) for i in xrange(*s.indices())]
def index_all(self, item):
return self.index(item, False)
def is_visible(self, i):
node = self._getnode(i, 1)
return node.annotation[0] == 1
def is_visible_item(self, item):
node = self._index_lookup(item)
return node.annotation[0] == 1
def insert_sequence_all(self, position, sequence, visibility):
node = self._getnode(position,1)
self._insert_sequence_leftofnode(node, sequence, visibility)
def insert_sequence_leftof(self, target, sequence, visibility):
node = self._index_lookup(target)
self._insert_sequence_leftofnode(node, sequence, visibility)
def _insert_sequence_leftofnode(self, node, sequence, visibility):
for i in xrange(len(sequence)):
v = sequence[i]
viz = visibility[i]
newnode = self._makenode()
newnode.annotation = (1 if viz else 0, 1)
newnode.value = v
self._tree.addleft(newnode, node)
self._index_assign(v, newnode)
# Skew, split, and decrease_level are the AA balancing functions, as described
# at http://en.wikipedia.org/wiki/AA_tree . They have been modified
# substantially here to (1) maintain bidirectional linking and (2) maintain
# monoid annotations.
def skew(node, op=None):
L = node.leftchild
if (L is not None) and node.level == L.level:
node.leftchild = L.rightchild
if node.leftchild is not None:
node.leftchild.parent = node
L.rightchild = node
L.parent = node.parent
node.parent = L
if L.parent is not None:
if L.parent.leftchild is node:
L.parent.leftchild = L
else:
assert L.parent.rightchild is node
L.parent.rightchild = L
if op is not None:
L.annotation = node.annotation
node.annotation = op(node.leftchild.annotation, node.rightchild.annotation)
assert L.annotation == op(L.leftchild.annotation, L.rightchild.annotation)
# This assertion is the condition of associativity, guaranteed for any
# valid monoid operation.
return L
else:
return node
def split(node, op=None):
R = node.rightchild
if ((R is not None) and
(R.rightchild is not None) and
(node.level == R.rightchild.level)):
node.rightchild = R.leftchild
node.rightchild.parent = node
R.leftchild = node
R.parent = node.parent
node.parent = R
R.level += 1
if R.parent is not None:
if R.parent.leftchild is node:
R.parent.leftchild = R
else:
assert R.parent.rightchild is node
R.parent.rightchild = R
if op is not None:
R.annotation = node.annotation
node.annotation = op(node.leftchild.annotation, node.rightchild.annotation)
assert R.annotation == op(R.leftchild.annotation, R.rightchild.annotation)
# This assertion is the condition of associativity, guaranteed for any
# valid monoid operation.
return R
else:
return node
def decrease_level(node):
# Decrease the level of node if necessary. Returns true if a modification
# was made.
target = min(node.leftchild.level, node.rightchild.level) + 1
if target < node.level:
node.level = target
if target < node.rightchild.level:
node.rightchild.level = target
return True
return False
class AAMonoidTree(MonoidTree):
makenode = AANode
def _update_add(self, node, sentinel=None):
"""node must be an internal node one level above the leaves, with
two leaves itself."""
node.level = 2
while node is not sentinel:
#oldval = node.annotation
node.annotation = self.op(node.leftchild.annotation, node.rightchild.annotation)
node = skew(node, self.op)
node = split(node, self.op)
if node.parent is None:
self.root = node
node = node.parent
def _update_del(self, node, sentinel=None):
while node is not sentinel:
#oldval = node.annotation
#oldlevel = node.level
node.annotation = self.op(node.leftchild.annotation, node.rightchild.annotation)
decrease_level(node)
node = skew(node, self.op)
node.rightchild = skew(node.rightchild, self.op)
if node.rightchild.rightchild is not None:
node.rightchild.rightchild = skew(node.rightchild.rightchild, self.op)
node = split(node, self.op)
node.rightchild = split(node.rightchild, self.op)
#if (oldval == node.annotation) and (oldlevel == node.level):
# #Nodes above this point will not have changed
# break
if node.parent is None:
self.root = node
node = node.parent
class AATreeList(TreeList):
_treetype = AAMonoidTree
class AATreeHideList(TreeHideList):
_treetype = AAMonoidTree