Heavier Ball Puzzle example in Python

heavier_ball_puzzle.py

from random import shuffle
from functools import reduce
n = int(input())
a = [ 1 for x in range(0,n-1)]
a.append(2)
shuffle(a)

if(n == 3):
  print('Heavier ball found in iteration 0')
  
else:
  for i in range(0,n // 2):
    remain = a.pop()
    index = int(len(a)/2)
    if(index == 0):
        print('Heavier ball found in iteration {}'.format(i+1))
        break
    group_a = a[:index]
    group_b = a[index:]
    sum_a = reduce(lambda a,b: a+b,group_a)
    sum_b = reduce(lambda a,b: a+b,group_b)
    if(sum_a == sum_b):
      print('Heavier ball found in iteration {}'.format(i+1))
      break
    elif(sum_a > sum_b):
        a = a[:index]
    else:
        a = a[index:]