Queries to count groups of N students possible having sum of ratings within given range

Last Updated : 01 Jun, 2021

Given integers N and K representing the number of batches and number of students in each batch respectively, and a 2D array ratings[][] of size N * K where each row has ratings for every K students and Q queries of type {a, b}. The task for each query is to count the number of groups of N students possible by selecting a student from each batch, such that sum of ratings in each group lies in the range [a, b] inclusively.

Examples:

Input: N = 2, K = 3, ratings[][]= { {1, 2, 3}, {4, 5, 6} }, Q = 2, Queries[][]={ {6, 6}, {1, 6} }
Output: 2 3
Explanation:
All possible groups of size N(=2) are:
1 + 4 = 5
1 + 5 = 6
1 + 6 = 7
2 + 4 = 6
2 + 5 = 7
2 + 6 = 8
3 + 4 = 7
3 + 5 = 8
3 + 6 = 9
Query 1: The groups whose sum in range of (6, 6) inclusive are (1 + 5), (2 + 4) is 2.
Query 2: The groups whose sum in range of (1, 6) inclusive are (1 + 4), (1 + 5), (2 + 4) is 3.
Input: N = 3, K = 3, ratings[][]={ {1, 2, 3}, {4, 5, 6}, {7, 8, 9} }, Q = 2, Queries[][]={ {10, 13}, {1, 7} }
Output: 4 0
Explanation:
Out of All possible groups of size N(=3):
Query 1: The groups whose sum in range (10, 13) inclusive is (1 + 4 + 7), (1 + 5 + 7), (2 + 4 + 7), (1 + 4 + 8) is 4.
Query 2: The groups whose sum in range of (1, 7) inclusive is 0.

Naive Approach: The simplest approach is to use recursion to generate all possible groups of size N. At each step of recursion calculate the sum that lies within the range of a query and find the number of groups that lie in the given range.

Time Complexity: O(N^K)
Auxiliary Space: O(1)

Efficient Approach: The above approach can be optimized by using Dynamic Programming. The idea is that if the number of times the sum S appears in the i^th row is known, then the Prefix Sum Technique can be used to answer all the queries in constant time. So in this way, the Overlapping Subproblems are calculated only once reducing the exponential time into polynomial time. Below are the steps:

Initialize auxiliary array dp[][] where dp[i][sum] is the number of times a sum is present in the i^th row.
For each batch, i iterate through all possible sum S, and for each j students, if sum S is greater than the current rating ratings[i][j] the update the current dp state dp[i][S] as:

dp[i][S] = dp[i][S] + dp[i – 1][sum – rating[i][j]]

After the above steps, dp[N – 1][sum], which is the number of times the sum appears in the (N – 1)^th row.
To answer each queries efficiently find the prefix sum of the last row i.e., dp[N – 1][S] for all values of sum S.
Now, the number of ways to form groups in the given range [a, b] is dp[N – 1][b] – dp[N – 1][a – 1].

Below is the implementation of the above approach:

C++

// C++ program for the above approach
 
#include <bits/stdc++.h>
using namespace std;
 
// Given n batches and k students
#define n 2
#define k 3
 
// Function to count number of
// ways to get given sum groups
void numWays(int ratings[n][k], int queries[][2])
{
 
    // Initialise dp array
    int dp[n][10000 + 2];
 
    // Mark all 1st row values as 1
    // since the mat[0][i] is all
    // possible sums in first row
    for (int i = 0; i < k; i++)
        dp[0][ratings[0][i]] += 1;
 
    // Fix the ith row
    for (int i = 1; i < n; i++) {
 
        // Fix the sum
        for (int sum = 0; sum <= 10000; sum++)
        {
            // Iterate through all
            // values of ith row
            for (int j = 0; j < k; j++)
            {
                // If sum can be obtained
                if (sum >= ratings[i][j])
                    dp[i][sum]
                        += dp[i - 1][sum - ratings[i][j]];
            }
        }
    }
 
    // Find the prefix sum of last row
    for (int sum = 1; sum <= 10000; sum++)
    {
        dp[n - 1][sum] += dp[n - 1][sum - 1];
    }
 
    // Traverse each query
 
    for (int q = 0; q < 2; q++)
    {
        int a = queries[q][0];
        int b = queries[q][1];
 
        // No of ways to form groups
        cout << dp[n - 1][b] - dp[n - 1][a - 1] << " ";
    }
}
 
// Driver Code
int main()
{
 
    // Given ratings
    int ratings[n][k] = { { 1, 2, 3 }, { 4, 5, 6 } };
 
    // Given Queries
    int queries[][2] = { { 6, 6 }, { 1, 6 } };
 
    // Function Call
    numWays(ratings, queries);
 
    return 0;
}

Java

// Java program for the above approach
 
import java.util.*;
public class Main {
 
    // Function to count number of
    // ways to get given sum groups
    public static void
    numWays(int[][] ratings,
            int queries[][],
            int n, int k)
    {
 
        // Initialise dp array
        int dp[][] = new int[n][10000 + 2];
 
        // Mark all 1st row values as 1
        // since the mat[0][i] is all
        // possible sums in first row
        for (int i = 0; i < k; i++)
            dp[0][ratings[0][i]] += 1;
 
        // Fix the ith row
        for (int i = 1; i < n; i++)
        {
 
            // Fix the sum
            for (int sum = 0; sum <= 10000; sum++)
            {
 
                // Iterate through all
                // values of ith row
                for (int j = 0; j < k; j++)
                {
 
                    // If sum can be obtained
                    if (sum >= ratings[i][j])
                        dp[i][sum]
                            += dp[i - 1]
                               [sum - ratings[i][j]];
                }
            }
        }
 
        // Find the prefix sum of last row
        for (int sum = 1; sum <= 10000; sum++) {
            dp[n - 1][sum] += dp[n - 1][sum - 1];
        }
 
        // Traverse each query
        for (int q = 0; q < queries.length; q++) {
 
            int a = queries[q][0];
            int b = queries[q][1];
 
            // No of ways to form groups
            System.out.print(dp[n - 1][b] - dp[n - 1][a - 1]
                             + " ");
        }
    }
 
    // Driver Code
    public static void main(String args[])
    {
        // Given N batches and K students
        int N = 2, K = 3;
 
        // Given ratings
        int ratings[][] = { { 1, 2, 3 }, { 4, 5, 6 } };
 
        // Given Queries
        int queries[][] = { { 6, 6 }, { 1, 6 } };
 
        // Function Call
        numWays(ratings, queries, N, K);
    }
}

Python3

# Python3 program for the
# above approach
 
# Function to count number of
# ways to get given sum groups
def numWays(ratings, queries,
            n, k):
   
    # Initialise dp array
    dp = [[0 for i in range(10002)]
             for j in range(n)];
 
    # Mark all 1st row values as 1
    # since the mat[0][i] is all
    # possible sums in first row
    for i in range(k):
        dp[0][ratings[0][i]] += 1;
 
    # Fix the ith row
    for i in range(1, n):
 
        # Fix the sum
        for sum in range(10001):
 
            # Iterate through all
            # values of ith row
            for j in range(k):
 
                # If sum can be obtained
                if (sum >= ratings[i][j]):
                    dp[i][sum] += dp[i - 1][sum -
                                            ratings[i][j]];
 
    # Find the prefix sum of
    # last row
    for sum in range(1, 10001):
        dp[n - 1][sum] += dp[n - 1][sum - 1];
 
    # Traverse each query
    for q in range(len(queries)):
        a = queries[q][0];
        b = queries[q][1];
 
        # No of ways to form groups
        print(dp[n - 1][b] -
              dp[n - 1][a - 1],
              end = " ");
 
# Driver Code
if __name__ == '__main__':
   
    # Given N batches and
    # K students
    N = 2;
    K = 3;
 
    # Given ratings
    ratings = [[1, 2, 3],
               [4, 5, 6]];
 
    queries = [[6, 6],
               [1, 6]];
 
    # Function Call
    numWays(ratings, queries, N, K);
 
# This code is contributed by 29AjayKumar

C#

// C# program for the above approach
using System;
 
class GFG{
 
// Function to count number of
// ways to get given sum groups
public static void numWays(int[,] ratings,
                           int [,]queries,
                           int n, int k)
{
     
    // Initialise dp array
    int [,]dp = new int[n, 10000 + 2];
 
    // Mark all 1st row values as 1
    // since the mat[0,i] is all
    // possible sums in first row
    for(int i = 0; i < k; i++)
        dp[0, ratings[0, i]] += 1;
 
    // Fix the ith row
    for(int i = 1; i < n; i++)
    {
         
        // Fix the sum
        for(int sum = 0; sum <= 10000; sum++)
        {
             
            // Iterate through all
            // values of ith row
            for(int j = 0; j < k; j++)
            {
 
                // If sum can be obtained
                if (sum >= ratings[i, j])
                    dp[i, sum] += dp[i - 1,
                                   sum - ratings[i, j]];
            }
        }
    }
 
    // Find the prefix sum of last row
    for(int sum = 1; sum <= 10000; sum++)
    {
        dp[n - 1, sum] += dp[n - 1, sum - 1];
    }
 
    // Traverse each query
    for(int q = 0; q < queries.GetLength(0); q++)
    {
        int a = queries[q, 0];
        int b = queries[q, 1];
 
        // No of ways to form groups
        Console.Write(dp[n - 1, b] -
                      dp[n - 1, a - 1] + " ");
    }
}
 
// Driver Code
public static void Main(String []args)
{
     
    // Given N batches and K students
    int N = 2, K = 3;
 
    // Given ratings
    int [,]ratings = { { 1, 2, 3 }, { 4, 5, 6 } };
 
    // Given Queries
    int [,]queries = { { 6, 6 }, { 1, 6 } };
 
    // Function Call
    numWays(ratings, queries, N, K);
}
}
 
// This code is contributed by Amit Katiyar

Javascript

<script>
 
// Javascript program for the above approach
 
// Given n batches and k students
var n = 2;
var k = 3;
 
// Function to count number of
// ways to get given sum groups
function numWays(ratings, queries)
{
     
    // Initialise dp array
    var dp = Array.from(
        Array(n), ()=>Array(10002).fill(0));
 
    // Mark all 1st row values as 1
    // since the mat[0][i] is all
    // possible sums in first row
    for(var i = 0; i < k; i++)
        dp[0][ratings[0][i]] += 1;
 
    // Fix the ith row
    for(var i = 1; i < n; i++)
    {
         
        // Fix the sum
        for(var sum = 0; sum <= 10000; sum++)
        {
             
            // Iterate through all
            // values of ith row
            for(var j = 0; j < k; j++)
            {
                 
                // If sum can be obtained
                if (sum >= ratings[i][j])
                    dp[i][sum]+= dp[i - 1][sum - ratings[i][j]];
            }
        }
    }
 
    // Find the prefix sum of last row
    for(var sum = 1; sum <= 10000; sum++)
    {
        dp[n - 1][sum] += dp[n - 1][sum - 1];
    }
 
    // Traverse each query
    for(var q = 0; q < 2; q++)
    {
        var a = queries[q][0];
        var b = queries[q][1];
 
        // No of ways to form groups
        document.write(dp[n - 1][b] -
                       dp[n - 1][a - 1] + " ");
    }
}
 
// Driver Code
 
// Given ratings
var ratings = [ [ 1, 2, 3 ], [ 4, 5, 6 ] ];
 
// Given Queries
var queries = [ [ 6, 6 ], [ 1, 6 ] ];
 
// Function Call
numWays(ratings, queries);
 
// This code is contributed by famously
 
</script>

Output

2 3

Time Complexity: O(N*maxSum*K), where maxSum is the maximum sum.
Auxiliary Space: O(N*maxSum), where maxSum is the maximum sum.

Attention reader! Don’t stop learning now. Get hold of all the important DSA concepts with the DSA Self Paced Course at a student-friendly price and become industry ready. To complete your preparation from learning a language to DS Algo and many more, please refer Complete Interview Preparation Course.

In case you wish to attend live classes with industry experts, please refer DSA Live Classes

My Personal Notes

Related Articles

C++

Java

Python3

C#

Javascript