forked from adnanaziz/EPIJudge
-
Notifications
You must be signed in to change notification settings - Fork 0
/
bonus.py
36 lines (29 loc) · 1.3 KB
/
bonus.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
import collections
import heapq
from test_framework import generic_test
def calculate_bonus(productivity):
# Stores (productivity, index)-pair in min_heap where ordered by
# productivity.
EmployeeData = collections.namedtuple('EmployeeData',
('productivity', 'index'))
min_heap = [EmployeeData(p, i) for i, p in enumerate(productivity)]
heapq.heapify(min_heap)
# Initially assigns one ticket to everyone.
tickets = [1] * len(productivity)
# Fills tickets from lowest rating to highest rating.
while min_heap:
next_dev = heapq.heappop(min_heap)[1]
# Handles the left neighbor.
if next_dev > 0 and productivity[next_dev] > productivity[next_dev
- 1]:
tickets[next_dev] = tickets[next_dev - 1] + 1
# Handles the right neighbor.
if (next_dev + 1 < len(tickets)
and productivity[next_dev] > productivity[next_dev + 1]):
tickets[next_dev] = max(tickets[next_dev],
tickets[next_dev + 1] + 1)
return sum(tickets)
if __name__ == '__main__':
exit(
generic_test.generic_test_main("bonus.py", 'bonus.tsv',
calculate_bonus))