Find out how many distinct triplets of numbers, from give N integers , have their sum divisible by integer M

Question

Given N integers. A triplet of numbers is a set of three numbers whose sum is divisible by a constant M. Find out how many distinct triplets of numbers, from given N integers , have their sum divisible by M. What would be an optimal algorithm to find the answer to this problem ?

Please Note: Every element is treated as different irrespective of their actual value .

This is the link to the actual problem

Answer 1

There's an O(N + M^2) solution.

First create an array C[M] where C[i] contains the number of the input numbers that are i mod M. This takes O(N) time.

Then for each i <= j <= k with i+j+k=0 mod M, count triples (a, b, c) with a=i (mod M), b=j (mod M), c=k(mod M).

We can do this efficiently by looping over i and j, and finding the unique k that causes the sum to be 0 (mod M), ignoring the combination if k<j. We have to be careful counting to make sure that we count unique triples; if i, j, k are distinct the number of triples is simply C[i] * C[j] * C[k], but if two or more of the indices are the same, we have to be careful not to allow the same number to be reused. For example, if i=j!=k, then there's choose(C[i], 2) * C[k] triples.

Here's complete code, which produces the correct answer 26 for the given test case:

def choose2(n):
    return n * (n-1) // 2

def choose3(n):
    return n * (n-1) * (n-2) // 6

def triples(A, M):
    C = [0] * M
    for a in A:
        C[a % M] += 1
    T = 0
    for i in range(M):
        for j in range(i, M):
            k = (-i - j) % M
            if k < j:
                continue
            if i == j == k:
                T += choose3(C[i])
            elif i == j:
                T += choose2(C[i]) * C[k]
            elif j == k:
                T += C[i] * choose2(C[j])
            else:
                T += C[i] * C[j] * C[k]
    return T

print(triples([1, 10, 4, 3, 2, 5, 0, 1, 9, 5], 5))