Advertisement
Guest User

Using SortedContainers Types for IndexedList

a guest
Dec 1st, 2014
222
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 2.07 KB | None | 0 0
  1. from sortedcontainers import SortedListWithKey, SortedDict
  2. from operator import itemgetter
  3.  
  4. class IndexedList(object):
  5.     def __init__(self, *args, **kwargs):
  6.         self.pk = 0
  7.         self.indices = {}
  8.         self.uniques = {'_pk': SortedDict()}
  9.         self.keys = {'_pk': itemgetter(0)}
  10.  
  11.     def append(self, item):
  12.         self.pk += 1
  13.         item = (self.pk, item)
  14.         for name in self.uniques:
  15.             self.uniques[name][self.keys[name](item)] = item
  16.         for name in self.indices:
  17.             self.indices[name].add(item)
  18.  
  19.     def __iter__(self):
  20.         return self.uniques['_pk'].values()
  21.  
  22.     def add_unique_index(self, name, key):
  23.         key = lambda pair: key(pair[1])
  24.         self.uniques[name] = SortedDict(zip(map(key, self), self))
  25.         self.keys[name] = key
  26.  
  27.     def add_index(self, name, key):
  28.         key = lambda pair: key(pair[1])
  29.         items = zip(map(key, self), self)
  30.         self.indices[name] = SortedListWithKey(items, key=itemgetter(0))
  31.  
  32.     def remove_unique(self, name, key):
  33.         item = self.uniques[name][key]
  34.         for name in self.uniques:
  35.             del self.uniques[name][self.keys[name](item)]
  36.         for name in self.indices:
  37.             self.indices[name].remove(item)
  38.  
  39.     def remove_index(self, name, key):
  40.         start = self.indices[name].bisect_left(key)
  41.         end = self.indices[name].bisect_right(key)
  42.         items = map(itemgetter(1), self.indices[name][start:end])
  43.         for item in items:
  44.             self.remove_unique('_pk', item[0])
  45.  
  46. book = IndexedList()
  47. book.add_unique_index('order_id', lambda o: o.id)
  48. book.add_index('account_id', lambda o: o.account_id)
  49. book.add_index('price_time', lambda o: (o.price, o.time))
  50.  
  51. def add_order(order):
  52.     book.append(order)
  53.  
  54. def del_order(order):
  55.     book.remove_unique('order_id', order.id)
  56.  
  57. def close_account(account_id):
  58.     book.remove_index('account_id', account_id)
  59.  
  60. def iter_matches(limit_price):
  61.     for (price, time), order in book.indices['price_time']:
  62.         if price > limit_price:
  63.             break
  64.         yield order
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement