blob: d2857993756df98460d5df8f2d4d2e37faa61ae7 (
plain)
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
|
#!/usr/bin/env python3
import sys
def equalize_array(ints: list, x: int, y: int) -> str:
score = 0
# Calculate the needed values
max_value = max(ints)
needed = [max_value - i for i in ints]
# If we have at least two items, and y is less than 2 × x, lets remove two
# values at a time while we can
if len(ints) > 1 and y < x * 2:
while True:
# Sort the positions by the values they hold
sorted_index = sorted(
range(len(ints)),
key=lambda index: needed[index],
reverse=True
)
# If we don't have two non zero values, we are done with level 2
if needed[sorted_index[1]] == 0:
break
# Take one off each number at that position
needed[sorted_index[0]] -= 1
needed[sorted_index[1]] -= 1
score += y
# As level one takes one off each number, we simply multiple the remaining
# needed values by the x value
score += sum(needed) * x
return score
def main():
# Convert input into integers
array = [int(n) for n in sys.argv[1:]]
y = array.pop()
x = array.pop()
result = equalize_array(array, x, y)
print(result)
if __name__ == '__main__':
main()
|