hacker-cup/2017.1-RoundOne/a.PieProgress.py

import numpy as np

def calcCost(sol, S):
    res = 0
    for i, n in enumerate(sol):
        for j in range(n):
            # print('res+=', S[i][j])
            res += S[i][j]
        if n > 0:
            # print('resb+=', pow(2, n), n)
            res += pow(n, 2)
    return res

def solveRec(N, M, S):
    if N == 0:
        return []
    else:
        sol = solveRec(N-1, M, S)
        sol.append(0)
        minCost = -1
        for i in range(N):
            if sol[i] >= len(S[i]):
                continue
            newSol = list(sol)
            newSol[i] += 1
            newCost = calcCost(newSol, S)
            if newCost < minCost or minCost == -1:
                minCost = newCost
                bestSol = newSol
        return bestSol

def solve(N, M, S):
    # print("N={0}, M={1}".format(N, M))
    # print(S)
    sol = solveRec(N,M,S)
    # print(sol)
    return calcCost(sol, S)

if __name__ == "__main__":
    import fileinput
    f = fileinput.input()

    T = int(f.readline())
    for case in range(1, T+1):
        N, M = [int(i) for i in f.readline().split()]
        S = [sorted([int(i) for i in f.readline().split()]) for x in range(N)]
        solution = solve(N, M, S)
        print("Case #{0}: {1}".format(case, solution))