1import bintrees
2class ClientsCreditsInfo:
3 def __init__(self):
4 self.offset = 0
5 self.clientToCredit = {}
6 self.creditToClient = bintrees.RBTree()
7
8 def insert(self, client_id, credit):
9 self.remove(client_id)
10 self.clientToCredit[client_id] = credit - self.offset
11 self.creditToClient.setDefault(credit - self.offset, set()).add(client_id)
12
13 def remove(self, client_id):
14 if client_id in self.clientToCredit:
15 credit = self.clientToCredit[client_id]
16 del self.clientToCredit[client_id]
17 self.creditToClient[credit].remove(client_id)
18 if not self.creditToClient[credit]: del self.creditToClient[credit]
19 return True
20 return False
21
22 def lookup(self, client_id):
23 if client_id in self.clientToCredit: return self.clientToCredit[client_id]
24 return False
25
26 def add_all(self, credit):
27 self.offset += credit
28
29 def max(self):
30 res = self.creditToClient.max_item()[1]
31 if res: return next(iter(res))
32 return False