forked from adnanaziz/EPIJudge
-
Notifications
You must be signed in to change notification settings - Fork 0
/
adding_credits.py
84 lines (64 loc) · 2.5 KB
/
adding_credits.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
import bintrees
from test_framework import generic_test
from test_framework.test_failure import TestFailure
class ClientsCreditsInfo:
def __init__(self):
self._offset = 0
self._client_to_credit = {}
self._credit_to_clients = bintrees.RBTree()
def insert(self, client_id: str, c: int) -> None:
self.remove(client_id)
self._client_to_credit[client_id] = c - self._offset
self._credit_to_clients.setdefault(c - self._offset,
set()).add(client_id)
def remove(self, client_id: str) -> bool:
credit = self._client_to_credit.get(client_id)
if credit is not None:
self._credit_to_clients[credit].remove(client_id)
if not self._credit_to_clients[credit]:
del self._credit_to_clients[credit]
del self._client_to_credit[client_id]
return True
return False
def lookup(self, client_id: str) -> int:
credit = self._client_to_credit.get(client_id)
return -1 if credit is None else credit + self._offset
def add_all(self, C: int) -> None:
self._offset += C
def max(self) -> str:
if not self._credit_to_clients:
return ''
clients = self._credit_to_clients.max_item()[1]
return '' if not clients else next(iter(clients))
def client_credits_info_tester(ops):
cr = ClientsCreditsInfo()
for x in ops:
op = x[0]
s_arg = x[1]
i_arg = x[2]
if op == 'ClientsCreditsInfo':
pass
if op == 'max':
result = cr.max()
if result != s_arg:
raise TestFailure('Max: return value mismatch')
elif op == 'remove':
result = cr.remove(s_arg)
if result != i_arg:
raise TestFailure('Remove: return value mismatch')
elif op == 'insert':
cr.insert(s_arg, i_arg)
elif op == 'add_all':
cr.add_all(i_arg)
elif op == 'lookup':
result = cr.lookup(s_arg)
if result != i_arg:
raise TestFailure('Lookup: return value mismatch')
def program_config(config):
config.analyze_complexity = False
if __name__ == '__main__':
exit(
generic_test.generic_test_main('adding_credits.py',
'adding_credits.tsv',
client_credits_info_tester,
program_config=program_config))