Number of triplets such that each value is less than N and each pair sum is a multiple of K

Given two integers N and K. Find the numbers of triplets (a, b, c) such that 0 ? a, b, c ? N and (a + b), (b + c) and (c + a) are multiples of K.

Examples:

Input: N = 3, K = 2 
Output: 9 
Triplets possible are: 
{(1, 1, 1), (1, 1, 3), (1, 3, 1) 
(1, 3, 3), (2, 2, 2), (3, 1, 1) 
(3, 1, 1), (3, 1, 3), (3, 3, 3)}

Input: N = 5, K = 3 
Output: 1 
Only possible triplet is (3, 3, 3) 

Approach: Given that (a + b), (b + c) and (c + a) are multiples of K. Hence, we can say that (a + b) % K = 0, (b + c) % K = 0 and (c + a) % K = 0. 
If a belongs to the x modulo class of K then b should be in the (K – x)th modulo class using the first condition. 
From the second condition, it can be seen that c belongs to the x modulo class of K. Now as both a and c belong to the same modulo class and they have to satisfy the third relation which is (a + c) % K = 0. It could be only possible if x = 0 or x = K / 2. 
When K is an odd integer, x = K / 2 is not valid. 

Hence to solve the problem, count the number of elements from 0 to N in the 0^th modulo class and the (K / 2)^th modulo class of K. 

• If K is odd then the result is cnt[0]³
• If K is even then the result is cnt[0]³ + cnt[K / 2]³.

Below is the implementation of the above approach:  

C++

// C++ implementation of the above approach
#include <bits/stdc++.h>
using namespace std;

// Function to return the number of triplets
int NoofTriplets(int N, int K)
{
    int cnt[K];

    // Initializing the count array
    memset(cnt, 0, sizeof(cnt));

    // Storing the frequency of each modulo class
    for (int i = 1; i <= N; i += 1) {
        cnt[i % K] += 1;
    }

    // If K is odd
    if (K & 1)
        return cnt[0] * cnt[0] * cnt[0];

    // If K is even
    else {
        return (cnt[0] * cnt[0] * cnt[0]
                + cnt[K / 2] * cnt[K / 2] * cnt[K / 2]);
    }
}

// Driver Code
int main()
{
    int N = 3, K = 2;

    // Function Call
    cout << NoofTriplets(N, K);

    return 0;
}

Java

// Java implementation of the approach
import java.util.Arrays;

class GFG
{

    // Function to return the number of triplets
    static int NoofTriplets(int N, int K)
    {
        int[] cnt = new int[K];

        // Initializing the count array
        Arrays.fill(cnt, 0, cnt.length, 0);

        // Storing the frequency of each modulo class
        for (int i = 1; i <= N; i += 1)
        {
            cnt[i % K] += 1;
        }

        // If K is odd
        if ((K & 1) != 0)
        {
            return cnt[0] * cnt[0] * cnt[0];
        }
        // If K is even
        else
        {
            return (cnt[0] * cnt[0] * cnt[0]
                    + cnt[K / 2] * cnt[K / 2] * cnt[K / 2]);
        }
    }

    // Driver Code
    public static void main(String[] args)
    {

        int N = 3, K = 2;

        // Function Call
        System.out.println(NoofTriplets(N, K));
    }
}

// This code is contributed by Princi Singh

C#

// C# implementation of the approach
using System;

class GFG
{

    // Function to return the number of triplets
    static int NoofTriplets(int N, int K)
    {
        int[] cnt = new int[K];

        // Initializing the count array
        Array.Fill(cnt, 0, cnt.Length, 0);

        // Storing the frequency of each modulo class
        for (int i = 1; i <= N; i += 1)
        {
            cnt[i % K] += 1;
        }

        // If K is odd
        if ((K & 1) != 0)
        {
            return cnt[0] * cnt[0] * cnt[0];
        }
        // If K is even
        else
        {
            return (cnt[0] * cnt[0] * cnt[0]
                    + cnt[K / 2] * cnt[K / 2] * cnt[K / 2]);
        }
    }

    // Driver Code
    static public void Main ()
    {
            int N = 3, K = 2;

        // Function Call
        Console.Write(NoofTriplets(N, K));
    }
}

// This code is contributed by ajit

Python3

# Python3 implementation of the above approach

# Function to return the number of triplets
def NoofTriplets(N, K) :

    # Initializing the count array
    cnt = [0]*K;

    # Storing the frequency of each modulo class
    for i in range(1, N + 1) :
        cnt[i % K] += 1;

    # If K is odd
    if (K & 1) :
        rslt = cnt[0] * cnt[0] * cnt[0];
        return rslt

    # If K is even
    else :
        rslt = (cnt[0] * cnt[0] * cnt[0] +
                cnt[K // 2] * cnt[K // 2] * cnt[K // 2]);
        return rslt

# Driver Code
if __name__ == "__main__" :

    N = 3; K = 2;

    # Function Call
    print(NoofTriplets(N, K));

# This code is contributed by AnkitRai01

Javascript

<script>

// Javascript implementation of the above approach

// Function to return the number of triplets
function NoofTriplets(N, K)
{
    let cnt = Array(K);

    for(let i = 0; i < K; i++)
        cnt[i] = 0;

    // Storing the frequency of
    // each modulo class
    for(let i = 1; i <= N; i += 1)
    {
        cnt[i % K] += 1;
    }

    // If K is odd
    if (K & 1)
        return cnt[0] * cnt[0] * cnt[0];

    // If K is even
    else
    {
        return (cnt[0] * cnt[0] * cnt[0] +
                 cnt[K / 2] * cnt[K / 2] *
                 cnt[K / 2]);
    }
}

// Driver Code
let N = 3;
let K = 2;

// Function Call
document.write(NoofTriplets(N, K));

// This code is contributed by mohit kumar 29

</script>

Output:

9

Time Complexity: O(N)

Auxiliary Space: O(K)