def calculate_medians_with_rbt(numbers):
"""Calculate the sum of all 10,000 medians, modulo 10000"""
smaller = RBTree() # For storing the smaller half of numbers
larger = RBTree() # For storing the larger half of numbers
medians = []
for number in numbers:
if not len(smaller) or smaller.max_item()[0] > number:
smaller.insert(number, None)
if len(smaller) > len(larger) + 1:
larger.insert(smaller.max_item()[0], None)
else:
larger.insert(number, None)
if len(larger) > len(smaller) + 1:
smaller.insert(larger.min_item()[0], None)
if len(smaller) >= len(larger):
medians.append(smaller.max_item()[0])
else:
medians.append(larger.min_item()[0])
return medians
def calculate_view_from_above(A):
class Endpoint:
def __init__(self, is_left, line):
self.is_left = is_left
self.line = line
def __lt__(self, other):
return self.value() < other.value()
def value(self):
return self.line.left if self.is_left else self.line.right
sorted_endpoints = sorted([Endpoint(True, a)
for a in A] + [Endpoint(False, a) for a in A])
result = []
prev_xaxis = sorted_endpoints[0].value()
prev = None
active_line_segments = RBTree()
for endpoint in sorted_endpoints:
if active_line_segments and prev_xaxis != endpoint.value():
active_segment = active_line_segments.max_item()[1]
if prev is None:
prev = LineSegement(prev_xaxis, endpoint.value(),
active_segment.color,
active_segment.height)
else:
if (prev.height == active_segment.height
and prev.color == active_segment.color
and prev_xaxis == prev.right):
prev = prev._replace(right=endpoint.value())
else:
result.append(prev)
prev = LineSegement(prev_xaxis, endpoint.value(),
active_segment.color,
active_segment.height)
prev_xaxis = endpoint.value()
if endpoint.is_left:
active_line_segments[endpoint.line.height] = endpoint.line
else:
del active_line_segments[endpoint.line.height]
return result + [prev] if prev else result
def solve(par):
n, days = par
urn = RBTree()
cost = 0
for d in days:
for bill in d:
try:
urn[bill] += 1
except:
urn[bill] = 1
if urn.count >= 2:
m1, v1 = urn.max_item()
m2, v2 = urn.min_item()
if v1 == 1:
urn.discard(m1)
else:
urn[m1] = v1 - 1
if v2 == 1:
urn.discard(m2)
else:
urn[m2] = v2 - 1
cost += m1 - m2
return cost
class OrderBook(WebsocketClient):
def __init__(self, product_id='BTC-USD', log_to=None):
super(OrderBook, self).__init__(products=product_id)
self._asks = RBTree()
self._bids = RBTree()
self._client = PublicClient()
self._sequence = -1
self._log_to = log_to
if self._log_to:
assert hasattr(self._log_to, 'write')
self._current_ticker = None
@property
def product_id(self):
''' Currently OrderBook only supports a single product even though it is stored as a list of products. '''
return self.products[0]
def on_open(self):
self._sequence = -1
print("-- Subscribed to OrderBook! --\n")
def on_close(self):
print("\n-- OrderBook Socket Closed! --")
def reset_book(self):
self._asks = RBTree()
self._bids = RBTree()
res = self._client.get_product_order_book(product_id=self.product_id, level=3)
for bid in res['bids']:
self.add({
'id': bid[2],
'side': 'buy',
'price': Decimal(bid[0]),
'size': Decimal(bid[1])
})
for ask in res['asks']:
self.add({
'id': ask[2],
'side': 'sell',
'price': Decimal(ask[0]),
'size': Decimal(ask[1])
})
self._sequence = res['sequence']
def on_message(self, message):
if self._log_to:
pickle.dump(message, self._log_to)
sequence = message['sequence']
if self._sequence == -1:
self.reset_book()
return
if sequence <= self._sequence:
# ignore older messages (e.g. before order book initialization from getProductOrderBook)
return
elif sequence > self._sequence + 1:
self.on_sequence_gap(self._sequence, sequence)
return
msg_type = message['type']
if msg_type == 'open':
self.add(message)
elif msg_type == 'done' and 'price' in message:
self.remove(message)
elif msg_type == 'match':
self.match(message)
self._current_ticker = message
elif msg_type == 'change':
self.change(message)
self._sequence = sequence
def on_sequence_gap(self, gap_start, gap_end):
self.reset_book()
print('Error: messages missing ({} - {}). Re-initializing book at sequence.'.format(
gap_start, gap_end, self._sequence))
def add(self, order):
order = {
'id': order.get('order_id') or order['id'],
'side': order['side'],
'price': Decimal(order['price']),
'size': Decimal(order.get('size') or order['remaining_size'])
}
if order['side'] == 'buy':
bids = self.get_bids(order['price'])
if bids is None:
bids = [order]
else:
bids.append(order)
self.set_bids(order['price'], bids)
else:
asks = self.get_asks(order['price'])
if asks is None:
asks = [order]
else:
asks.append(order)
self.set_asks(order['price'], asks)
def remove(self, order):
price = Decimal(order['price'])
if order['side'] == 'buy':
bids = self.get_bids(price)
if bids is not None:
bids = [o for o in bids if o['id'] != order['order_id']]
if len(bids) > 0:
self.set_bids(price, bids)
else:
self.remove_bids(price)
else:
asks = self.get_asks(price)
if asks is not None:
asks = [o for o in asks if o['id'] != order['order_id']]
if len(asks) > 0:
self.set_asks(price, asks)
else:
self.remove_asks(price)
def match(self, order):
size = Decimal(order['size'])
price = Decimal(order['price'])
if order['side'] == 'buy':
bids = self.get_bids(price)
if not bids:
return
assert bids[0]['id'] == order['maker_order_id']
if bids[0]['size'] == size:
self.set_bids(price, bids[1:])
else:
bids[0]['size'] -= size
self.set_bids(price, bids)
else:
asks = self.get_asks(price)
if not asks:
return
assert asks[0]['id'] == order['maker_order_id']
if asks[0]['size'] == size:
self.set_asks(price, asks[1:])
else:
asks[0]['size'] -= size
self.set_asks(price, asks)
def change(self, order):
try:
new_size = Decimal(order['new_size'])
except KeyError:
return
try:
price = Decimal(order['price'])
except KeyError:
return
if order['side'] == 'buy':
bids = self.get_bids(price)
if bids is None or not any(o['id'] == order['order_id'] for o in bids):
return
index = [b['id'] for b in bids].index(order['order_id'])
bids[index]['size'] = new_size
self.set_bids(price, bids)
else:
asks = self.get_asks(price)
if asks is None or not any(o['id'] == order['order_id'] for o in asks):
return
index = [a['id'] for a in asks].index(order['order_id'])
asks[index]['size'] = new_size
self.set_asks(price, asks)
tree = self._asks if order['side'] == 'sell' else self._bids
node = tree.get(price)
if node is None or not any(o['id'] == order['order_id'] for o in node):
return
def get_current_ticker(self):
return self._current_ticker
def get_current_book(self):
result = {
'sequence': self._sequence,
'asks': [],
'bids': [],
}
for ask in self._asks:
try:
# There can be a race condition here, where a price point is removed
# between these two ops
this_ask = self._asks[ask]
except KeyError:
continue
for order in this_ask:
result['asks'].append([order['price'], order['size'], order['id']])
for bid in self._bids:
try:
# There can be a race condition here, where a price point is removed
# between these two ops
this_bid = self._bids[bid]
except KeyError:
continue
for order in this_bid:
result['bids'].append([order['price'], order['size'], order['id']])
return result
def get_ask(self):
return self._asks.min_key()
def get_asks(self, price):
return self._asks.get(price)
def remove_asks(self, price):
self._asks.remove(price)
def set_asks(self, price, asks):
self._asks.insert(price, asks)
def get_bid(self):
return self._bids.max_key()
def get_bids(self, price):
return self._bids.get(price)
def remove_bids(self, price):
self._bids.remove(price)
def set_bids(self, price, bids):
self._bids.insert(price, bids)
def get_ask_size(self):
return sum(order['size'] for order in self._asks.min_item()[1])
def get_bid_size(self):
return sum(order['size'] for order in self._bids.max_item()[1])
