60 lines
1.3 KiB
Python
60 lines
1.3 KiB
Python
def solution(pegs):
|
|
a, b = -1, -1
|
|
odd, even = sum_odd_even(pegs)
|
|
|
|
if len(pegs) % 2 == 0:
|
|
s = 2 * odd - 2 * even - pegs[-1] + pegs[0]
|
|
if s % 3 == 0:
|
|
a, b = 2*s/3, 1
|
|
else:
|
|
a, b = 2*s, 3
|
|
else:
|
|
s = 2 * odd - 2 * even + pegs[-1] + pegs[0]
|
|
a, b = 2*s, 1
|
|
|
|
if check_all_gear_size(a, b, pegs):
|
|
return [a, b]
|
|
return [-1, -1]
|
|
|
|
def sum_odd_even(pegs):
|
|
odd, even = 0, 0
|
|
for i, peg in enumerate(pegs):
|
|
if i % 2 == 1:
|
|
odd += peg
|
|
else:
|
|
even += peg
|
|
return odd, even
|
|
|
|
def check_all_gear_size(a, b, pegs):
|
|
if a < 2*b:
|
|
return False
|
|
|
|
p = list(pegs) if b == 1 else [b*p for p in pegs]
|
|
|
|
last_peg = p[0]
|
|
last_gear = a
|
|
for i in range(1, len(p)):
|
|
peg = p[i]
|
|
gap = peg - last_peg
|
|
if gap - last_gear < b:
|
|
return False
|
|
last_gear = gap - last_gear
|
|
last_peg = peg
|
|
|
|
return True
|
|
|
|
|
|
tests = [
|
|
([4, 30, 50], [12, 1]),
|
|
([4, 17, 50], [-1, -1]),
|
|
([1, 2], [-1, -1]),
|
|
([1, 3], [-1, -1]),
|
|
([1, 4], [2, 1]),
|
|
([1, 5], [8, 3]),
|
|
([1, 2, 3], [-1, -1]),
|
|
]
|
|
|
|
for i, o in tests:
|
|
result = solution(i)
|
|
print (i, result == o, result, o)
|