fork download
  1. # your code goes here
  2. def solve(arr):
  3. n = len(arr)
  4. presum = [0 for i in range(n)]
  5. presum[0] = arr[0]
  6. for i in range(1,n):
  7. presum[i] = presum[i-1] + arr[i]
  8.  
  9. sufsum = [0 for i in range(n)]
  10. sufsum[n-1] = arr[n-1]
  11. for i in range(n-2,-1,-1):
  12. sufsum[i] = sufsum[i+1] + arr[i]
  13.  
  14. minindex = 0
  15. mindiff = presum[n-1] ## math.inf
  16. for i in range(n-1):
  17. leftavg = presum[i] // (i+1)
  18. rightavg = sufsum[i+1] // (n-i-1)
  19. currdiff = abs(leftavg - rightavg)
  20. print(i, leftavg, rightavg, currdiff)
  21. if currdiff < mindiff:
  22. minindex = i
  23. mindiff = currdiff
  24.  
  25. return minindex+1
  26.  
  27. print(solve([2,1]))
Success #stdin #stdout 0.03s 8932KB
stdin
Standard input is empty
stdout
0 2 1 1
1