#!/usr/bin/python2 ## You are given a list of "pegs" with their location (positive integers, asc order). ## these pegs can accept gears with radius > 1 ## the goal is the last gear to be spinning twice as fast as the 1st one for any given list ## if it's possible return the first gear's radius in a list [a,b] where the ratio is a/b ## if not return [-1,-1] from fractions import gcd from functools import reduce # for tests from random import seed from random import random def simplify(ratio_nums): return list(map(reduce(gcd, ratio_nums).__rfloordiv__, ratio_nums)) def solution(pegs): if len(pegs) == 2: r1_by_three = 2 * (pegs[1] - pegs[0]) nums = [r1_by_three, 3] return simplify(nums) S = pegs[-1] - pegs[0] if len(pegs) % 2 == 0: C = 0 for i in range (1, len(pegs) - 1): C += (-1)**i * pegs[i] r1_by_three = 2 * S - 4 * C r = [r1_by_three/3] for i in range(len(pegs) -1): r.append(pegs[i+1]-pegs[i]-r[i]) if r[-1] < 1: return [-1, -1] if (r[0]) < 1 or (r[0]) > pegs[1]-pegs[0]-1: return [-1, -1] return simplify([r1_by_three, 3]) else: K = 0 for i in range (1, len(pegs) - 2): K += (-1)**i * pegs[i] r1 = 2*S - 4*K - 4*pegs[-1] + 4*pegs[-2] if r1 < 1: return [-1, -1] return [r1, 1] # test cases provided print(solution([4,30,50])) # [12, 1] print(solution([4,17,50])) # [-1, -1] # tests seed(1) def ran(): return int(random() * 1000) for i in range(1,1000): a=[ran(),ran(),ran(),ran(),ran(),ran()] a.sort() print(a) print(solution(a)) print("")