Can someone give an Optimized approach?

Question

Find distinct triplet (i,j,k) such that i<j<k and sum of the arr[i]+arr[j]+arr[k] is divisible by d?

If arr = [3, 3, 4, 7, 8] and d = 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<=n<=10^3
• 1<=arr[i]<=10^9

    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^3) approach.

Answer 1

In general with questions like this, it can be helpful to sort your array so that, at each index, you can assert some statement on vast sections. Consider sorting arr in terms of mod d. What I mean by this is the first section of the array will constitute all numbers x in arr such that x mod d = 0.

Therefore if you take the modulus with respect to d of every value you would get something like:

[0,...,0,1,...,1,......,d-1,...,d-1]

Now you simply have to find 3 values which sum to 0, d or 2d. This can be done in O(n^2).

Pseudocode implementation that finds a triplet summing to d, simply set d = 2d and run this again to find a triplet summing to 2d.

for i in arr:
    j = i + 1
    k = arr.length() - 1
    while (j < k):
        currSum = arr[i] + arr[j] + arr[k]
        if (currSum = d):
            return (i, j, k)
        else if (currSum > d):
            k--
        else j++

As Thierry Lathuille and anatoylg point out, this pseudocode implementation requires that you had followed the process of taking the modulus of every value resulting in the array I had displayed above. Furthermore we must only check 0, d and 2d as the largest triple sum possible is 3d-3 < 3d.