I’ve tried to solve this with brute force which gives me MLE(memory limit exceeded) for test set 2… Can someone suggest me a better approach for this…
Link to the question: Kick Start - Google’s Coding Competitions
My solution:
from itertools import combinations
def balls_in_boxes(n, m):
“”"Generate combinations of n balls in m boxes.
>>> list(balls_in_boxes(4, 2))
[(0, 4), (1, 3), (2, 2), (3, 1), (4, 0)]
>>> list(balls_in_boxes(3, 3))
[(0, 0, 3), (0, 1, 2), (0, 2, 1), (0, 3, 0), (1, 0, 2), (1, 1, 1), (1, 2, 0), (2, 0, 1), (2, 1, 0), (3, 0, 0)]
"""
for c in combinations(range(n + m - 1), m - 1):
yield tuple(b - a - 1 for a, b in zip((-1,) + c, c + (n + m - 1,)))
tc = int(input())
for t in range(0, tc):
n,k,p = map(int, input().split())
arr = [[0 for i in range(k)] for j in range(n)]
for i in range(0,n):
arr[i] = list(map(int, input().split()))
for i in range(0,n):
for j in range(1,k):
arr[i][j] = arr[i][j-1]+arr[i][j]
li = list(balls_in_boxes(p,n))
maxCost = 0
for tup in li:
cost = 0
#print(tup)
i = 0
for val in tup:
if(val == 0):
i += 1
continue
if(val > k):
break
else:
cost += arr[i][val-1]
i += 1
if cost > maxCost:
maxCost = cost
result = "Case #"+str(t+1)+": "+str(maxCost)
print(result)
