using sparse DP approach

k,m = map(int, input().split())
inp_list = []

for _ in range(k):
    _, *nums = map(int, input().split())
    inp_list.append([x * x % m for x in nums])

curr_results = {0}

for li in inp_list:
    new_results = set()
    for res in curr_results:
        for ele in li:
            new_results.add((res + ele) % m)
    curr_results = new_results
print(max(curr_results))

I hit this problem where I get 766 but expected output is 763, even with some of the answers here. Anyone else had this problem?

6 767 indata = ['2 488512261 423332742', '2 625040505 443232774', '1 4553600', '4 92134264 617699202 124100179 337650738', '2 778493847 932097163', '5 489894997 496724555 693361712 935903331 518538304']

I think the important python learning here is that while the docs only give product() examples with two iterables, product() can take an arbitrary list of iterables to produce all combinations. E.g., modifying an example from the docs:

list(product('ABCD', 'xy', '123'))
[('A', 'x', '1'), ('A', 'x', '2'), ('A', 'x', '3'), ('A', 'y', '1'), ('A', 'y', '2'), ('A', 'y', '3'), ('B', 'x', '1'), ('B', 'x', '2'), ('B', 'x', '3'), ('B', 'y', '1'), ('B', 'y', '2'), ('B', 'y', '3'), ('C', 'x', '1'), ('C', 'x', '2'), ('C', 'x', '3'), ('C', 'y', '1'), ('C', 'y', '2'), ('C', 'y', '3'), ('D', 'x', '1'), ('D', 'x', '2'), ('D', 'x', '3'), ('D', 'y', '1'), ('D', 'y', '2'), ('D', 'y', '3')]

This then lets us create the things to sum over and test for max all at once.

It should be noted that product() isn't magical with respect to run time, this is still a polynomial algorithm in O(n^K).

from itertools import product if name=='main': K,M=input().split() K=int(K) M=int(M) input_list=[] for i in range(K): N=list(map(int,input().split())) input_list.append(N) for i in range(len(input_list)): for j in range (len(input_list[i])): if(j==0): input_list[i].remove(input_list[i][j]) result=list(product(*input_list)) added=[] for i in range(len(result)): sum=0 for j in range(K): sum=sum+(result[i][j]**2) function=sum%M added.append(function) print(max(added))