Finding a distinct triplet where the sum of the values is divisible by a given number

Question

Given an array a and a number d, I want to count the number of distinct triplets (i,j,k) such that i<j<k and the sum aᵢ + aⱼ + aₖ is divisible by d.

Example: when a is [3, 3, 4, 7, 8] and d is 5 it should give count as 3:

• triplet with indices (1,2,3) 3+3+4=10
• triplet with indices (1,3,5) 3+4+8=15
• triplet with indices (2,3,4) 3+4+8=15

Constraints:

• 3 ≤ len(a) ≤ 10³
• 1 ≤ aᵢ ≤ 10⁹

def count_triplets(arr, d):
    n = len(arr)
    count = 0

    for i in range(n - 2):
        for j in range(i + 1, n - 1):
            for k in range(j + 1, n):
                if (arr[i] + arr[j] + arr[k]) % d == 0:
                     count += 1

    return count

I got this for my OA but only came up with O(n³) approach.

Answer 1

We can preprocess the array by setting each a_i to a_i modulo d.

---

We're asked for a sum of three numbers a_i + a_j + a_k congruent to 0 modulo d. Well gosh, that's a hard problem! What would be an easier one?

Suppose that a_i is already fixed, we can't change it. Well now we only have two numbers to worry about. What must they sum to? a_j + a_k has to be -a_i modulo d.

Hmmm, still don't know the answer, do we? But we're getting close. Fix a_j, and now we need merely pick suitable a_k. Scan and declare a match if a_k is -a_i - a_j modulo d, FTW.

---

Create a sorted index, with linear cost:

    pairs = sorted((val, i)  for i, val in enumerate(arr))

Now finding a desired a_k costs log2(len(arr)).

This is a good match for the sorted containers library. Total cost will still be O(N ** 2 × log(N)).

---

If we're willing to spend quadratic storage and time, I suppose that sorting (sum2, j, k) tuples, where sum2 = (a_j + a_k) % d, might be helpful. That knocks the scanning cost down to O(N × log(N ** 2)), but alas the quadratic cost of creating those tuples would dominate.

---

Hmmm, no need to sort. We can allocate a list of length d where each element is a list of tuples. That way we just ask for a_k in O(1) constant time, without need of binary search.

Answer 2

One of the weaknesses of the code is that we're actually finding every triple. We're not required to do so - the only thing that matters is the count of results.

That means that we can group the numbers into buckets modulo d (using a collections.Counter) and then work with the counts rather than with the individual array entries.

Given the values in the problem statement (a = [3, 3, 4, 7, 8] and d = 5), we get counts of 1✕2, 3✕3 and 1✕4.

Once we have the counts, then we find the ways those modular values can be combined to add (modulo d) to zero). In this case, the only successful combination is 3+3+4; there's 3 ways to pick two values from the three 3's, and only one way to pick a single 4, giving 3✕1 = 3 as the final count.

It's probably useful to use an O(d²) approach to find sums of pairs modulo-d; for each pair, that sum determines which bucket to examine for the third member of each triplet. You may find math.comb() useful when a single bucket contributes two or all three values.

Answer 3

Its possible to solve the problem in O(N + D*log(D)) time, using the Fast Fourier Transform (FFT). In a nutshell- reduce each value mod D as before, encode each value as x raised to that power; the sum of these values will give you a single polynomial of degree D-1 at most. Multiplying this polynomial with itself twice will give you a new polynomial where the coefficient of the x^a term indicates how many ways you can sum the original values (with repetition and without order) to get a.

This means if we sum the coefficients of x^D and x^(2D) terms, we can get something close to our desired answer, just requiring some cleanup to account for uniqueness and order. We can use the FFT to perform the polynomial multiplication. In fact, we don't even need to use the IFFT (inverse FFT) to convert back- due to the math plays out of adding B powers of a Bth root of unity for B terms, all but the k*Bth order terms will just cancel out when summed. Simply put, this means we can sum the values of the FFT and just divide by the number of terms to get the count we want (just without uniqueness / order).

To get uniqueness / order, we just need to do some dirty work using some inclusion / exclusion logic to discard elements that were double-counted, and add back only unique triplets where repetition in the original array allows. To deal with the ordering constraint, we merely divide the total count by 6 (3!).

To demonstrate, here is working code that utilizes this approach:

from collections import Counter
from numpy.fft import fft


def count_triplets(target, nums):
    multiplicity = Counter(nums)
    
    counter = Counter(x % target for x in nums)
    poly = [counter[x] for x in range(target)]
    
    # get counts without ordering nor uniqueness criteria
    count = (fft(poly) ** 3).sum() / target
    
    # use inclusion / exclusion to establish uniqueness
    for x, c in multiplicity.items():
        # consider each (x, x, y) such that x + x + y = 0 (mod target)
        
        # count the number of such y's
        val = poly[-2 * x % target]
        
        # handle the case where x = y, including for duplicate x's
        if 3 * x % target == 0:
            val -= c
            count -= c ** 3
        
        # remove these triplets from our count unconditionally
        count -= val * 3 * c ** 2

        # handle the cases where (x, x, x) is entirely valid-  add unique ones back
        if c >= 3:
            val *= c * (c - 1) // 2
            if 3 * x % target == 0:
                count += c * (c - 1) * (c - 2)
        
        # add back rest of unique + valid (x, x, y) for x != y
        if c >= 2:
            count += val * 6
        
    # divide out equivalent permutations (establish order)
    count /= 6

    # the FFT introduces a very small amount of floating point imprecision 
    return round(count.real)

I've verified this solution against your provided code for thousands of test cases, so I'm fairly sure that its correct.

Answer 4

One improvement might be derived from changing every array entry to its modulo 5 value:

[3, 3, 4, 7, 8]
[3, 3, 4, 2, 3] has the same solution, or signed:
[-2, -2, -1, 2, -2]

Any solution containing a 3 may use indices 1, 2 or 5 arbitrarily.

A triplet giving a modulo 5 may be listed:

0 0 0 with 3 0s
...
4 4 2 with 2 4s and 1 2.

And we a searching with

1 2, 3 3s, 1 4

And notice that the last digit can be calculated to give 0 modulo 5.

So certainly a more optimal solution may be done. I just do not want to spoil your puzzling. (And some administration is needed, and I am more a Java guy.)

Answer 5

I'd prefer divisible_triplets(integers, divisor).
And it deserves a docstring.

You state I want to count the number of … -
I don't want to be caught telling people they want different.

To determine the count of such triplets, what if, in "the middle loop",
we knew just how many items of complementary residue followed after j?
Accumulation would be enough.
To establish that count, walk integers back to front, keeping
- one running count of items to follow per each residue 0 to divisor
- for each index, keep the counts for each j

When using both indices and values, consider enumerate().

def divisible_triplets(integers, divisor):
    """ Return count of ordered index triplets into integers
        such that the sum of the values indexed is divisible by divisor. """
    counts = [0] * divisor
    complements = [list(counts)]
    for i in reversed(integers[1:]):
        counts[-i%divisor] += 1                # complement in linear part
        complements.append(list(counts))
    complements = list(reversed(complements))  # complement once
    divisible_triplets = 0
    for i, ai in enumerate(integers[:-2]):
        for j, aj in enumerate(integers[i+1:-1], i+1):
            # print('i:', i, 'j:', j, complements[j])
            divisible_triplets += complements[j][(ai+aj)%divisor]
    return divisible_triplets

- should be O(n²).
(For an O(n+d²) approach, see Toby Speight's answer
 - trying to spell that out, I found no way to use itertools to spew out combinations with repetition limited by a Counter's counts.)