Number of ways to choose elements from the array such that their average is K - GeeksforGeeks

03 Mar, 2022
Given an array arr[] of N integers and an integer K. The task is to find the number of ways to select one or more elements from the array such that the average of the selected integers is equal to the given number K. 
Examples:  
Input: arr[] = {7, 9, 8, 9}, K = 8 Output: 5 {8}, {7, 9}, {7, 9}, {7, 8, 9} and {7, 8, 9}Input: arr[] = {3, 6, 2, 8, 7, 6, 5, 9}, K = 5 Output: 19Input: arr[] = {6, 6, 9}, K = 8 Output: 0  
Simple Approach: A simple solution would be to try for all possibilities since N can be large. Time complexity can be 2N.Efficient Approach: The above approach can be optimized by using dynamic programming to solve this problem. Suppose we are at the i_th index and let val be the current value of that index. We have two possibilities either to choose that element in the answer or discard the element. Hence, we are done now. We will also keep track of the number of elements in our current set of chosen elements. 
Following is the recursive formula. 
ways(index, sum, count) 
      = ways(index - 1, sum, count) 
        + ways(index - 1, sum + arr[index], count + 1)
Below is the implementation of the above approach:  
C++
Java
Python3
C#
Javascript
#include <bits/stdc++.h>using namespace std; #define MAX_INDEX 51#define MAX_SUM 2505 // This dp array is used to store our values// so that we don't have to calculate same// values again and againint dp[MAX_INDEX][MAX_SUM][MAX_INDEX]; int waysutil(int index, int sum, int count,             vector<int>& arr, int K){    // Base cases    // Index can't be less than 0    if (index < 0)        return 0;     if (index == 0) {         // No element is picked hence        // average cannot be calculated        if (count == 0)            return 0;        int remainder = sum % count;         // If remainder is non zero, we cannot        // divide the sum by count i.e. the average        // will not be an integer        if (remainder != 0)            return 0;        int average = sum / count;         // If we find an average return 1        if (average == K)            return 1;    }     // If we have already calculated this function    // simply return it instead of calculating it again    if (dp[index][sum][count] != -1)        return dp[index][sum][count];     // If we don't pick the current element    // simple recur for index -1    int dontpick = waysutil(index - 1,                            sum, count, arr, K);     // If we pick the current element add it to    // our current sum and increment count by 1    int pick = waysutil(index - 1,                        sum + arr[index],                        count + 1, arr, K);    int total = pick + dontpick;     // Store the value for the current function    dp[index][sum][count] = total;    return total;} // Function to return the number of waysint ways(int N, int K, int* arr){    vector<int> Arr;     // Push -1 at the beginning to    // make it 1-based indexing    Arr.push_back(-1);    for (int i = 0; i < N; ++i) {        Arr.push_back(arr[i]);    }     // Initialize dp array by -1    memset(dp, -1, sizeof dp);     // Call recursive function    // waysutil to calculate total ways    int answer = waysutil(N, 0, 0, Arr, K);    return answer;} // Driver codeint main(){    int arr[] = { 3, 6, 2, 8, 7, 6, 5, 9 };    int N = sizeof(arr) / sizeof(arr[0]);    int K = 5;    cout << ways(N, K, arr);     return 0;}
// Java implementation of the above approachimport java.util.*; class GFG{     static int MAX_INDEX = 51;    static int MAX_SUM = 2505;     // This dp array is used to store our values    // so that we don't have to calculate same    // values again and again    static int[][][] dp = new int[MAX_INDEX][MAX_SUM][MAX_INDEX];     static int waysutil(int index, int sum, int count,                        Vector<Integer> arr, int K)    {        // Base cases        // Index can't be less than 0        if (index < 0)        {            return 0;        }         if (index == 0)        {             // No element is picked hence            // average cannot be calculated            if (count == 0)            {                return 0;            }            int remainder = sum % count;             // If remainder is non zero, we cannot            // divide the sum by count i.e. the average            // will not be an integer            if (remainder != 0)            {                return 0;            }            int average = sum / count;             // If we find an average return 1            if (average == K)            {                return 1;            }        }         // If we have already calculated this function        // simply return it instead of calculating it again        if (dp[index][sum][count] != -1)        {            return dp[index][sum][count];        }         // If we don't pick the current element        // simple recur for index -1        int dontpick = waysutil(index - 1,                sum, count, arr, K);         // If we pick the current element add it to        // our current sum and increment count by 1        int pick = waysutil(index - 1,                sum + arr.get(index),                count + 1, arr, K);        int total = pick + dontpick;         // Store the value for the current function        dp[index][sum][count] = total;        return total;    }     // Function to return the number of ways    static int ways(int N, int K, int[] arr)    {        Vector<Integer> Arr = new Vector<>();         // Push -1 at the beginning to        // make it 1-based indexing        Arr.add(-1);        for (int i = 0; i < N; ++i)        {            Arr.add(arr[i]);        }         // Initialize dp array by -1        for (int i = 0; i < MAX_INDEX; i++)        {            for (int j = 0; j < MAX_SUM; j++)            {                for (int l = 0; l < MAX_INDEX; l++)                {                    dp[i][j][l] = -1;                }            }        }         // Call recursive function        // waysutil to calculate total ways        int answer = waysutil(N, 0, 0, Arr, K);        return answer;    }     // Driver code    public static void main(String args[])    {        int arr[] = {3, 6, 2, 8, 7, 6, 5, 9};        int N = arr.length;        int K = 5;        System.out.println(ways(N, K, arr));    }} /* This code contributed by PrinciRaj1992 */
# Python implementation of above approachimport numpy as np MAX_INDEX = 51MAX_SUM = 2505 # This dp array is used to store our values# so that we don't have to calculate same# values again and again # Initialize dp array by -1dp = np.ones((MAX_INDEX,MAX_SUM,MAX_INDEX)) * -1; def waysutil(index, sum, count, arr, K) :     # Base cases    # Index can't be less than 0    if (index < 0) :        return 0;     if (index == 0) :         # No element is picked hence        # average cannot be calculated        if (count == 0) :            return 0;                     remainder = sum % count;         # If remainder is non zero, we cannot        # divide the sum by count i.e. the average        # will not be an integer        if (remainder != 0) :            return 0;                     average = sum // count;         # If we find an average return 1        if (average == K) :            return 1;     # If we have already calculated this function    # simply return it instead of calculating it again    if (dp[index][sum][count] != -1) :        return dp[index][sum][count];     # If we don't pick the current element    # simple recur for index -1    dontpick = waysutil(index - 1,                            sum, count, arr, K);     # If we pick the current element add it to    # our current sum and increment count by 1    pick = waysutil(index - 1,                        sum + arr[index],                        count + 1, arr, K);                             total = pick + dontpick;     # Store the value for the current function    dp[index][sum][count] = total;         return total;  # Function to return the number of waysdef ways(N, K, arr) :     Arr = [];     # Push -1 at the beginning to    # make it 1-based indexing    Arr.append(-1);    for i in range(N) :        Arr.append(arr[i]);     # Call recursive function    # waysutil to calculate total ways    answer = waysutil(N, 0, 0, Arr, K);    return answer;  # Driver codeif __name__ == "__main__" :     arr = [ 3, 6, 2, 8, 7, 6, 5, 9 ];    N =len(arr);    K = 5;    print(ways(N, K, arr)); # This code is contributed by AnkitRai01
// C# implementation of the above approachusing System;using System.Collections.Generic;     class GFG{     static int MAX_INDEX = 51;    static int MAX_SUM = 2505;     // This dp array is used to store our values    // so that we don't have to calculate same    // values again and again    static int[,,] dp = new int[MAX_INDEX, MAX_SUM, MAX_INDEX];     static int waysutil(int index, int sum, int count,                        List<int> arr, int K)    {        // Base cases        // Index can't be less than 0        if (index < 0)        {            return 0;        }         if (index == 0)        {             // No element is picked hence            // average cannot be calculated            if (count == 0)            {                return 0;            }            int remainder = sum % count;             // If remainder is non zero, we cannot            // divide the sum by count i.e. the average            // will not be an integer            if (remainder != 0)            {                return 0;            }            int average = sum / count;             // If we find an average return 1            if (average == K)            {                return 1;            }        }         // If we have already calculated this function        // simply return it instead of calculating it again        if (dp[index,sum,count] != -1)        {            return dp[index, sum, count];        }         // If we don't pick the current element        // simple recur for index -1        int dontpick = waysutil(index - 1,                sum, count, arr, K);         // If we pick the current element add it to        // our current sum and increment count by 1        int pick = waysutil(index - 1,                sum + arr[index],                count + 1, arr, K);        int total = pick + dontpick;         // Store the value for the current function        dp[index,sum,count] = total;        return total;    }     // Function to return the number of ways    static int ways(int N, int K, int[] arr)    {        List<int> Arr = new List<int>();         // Push -1 at the beginning to        // make it 1-based indexing        Arr.Add(-1);        for (int i = 0; i < N; ++i)        {            Arr.Add(arr[i]);        }         // Initialize dp array by -1        for (int i = 0; i < MAX_INDEX; i++)        {            for (int j = 0; j < MAX_SUM; j++)            {                for (int l = 0; l < MAX_INDEX; l++)                {                    dp[i, j, l] = -1;                }            }        }         // Call recursive function        // waysutil to calculate total ways        int answer = waysutil(N, 0, 0, Arr, K);        return answer;    }     // Driver code    public static void Main(String []args)    {        int []arr = {3, 6, 2, 8, 7, 6, 5, 9};        int N = arr.Length;        int K = 5;        Console.WriteLine(ways(N, K, arr));    }} // This code is contributed by Princi Singh
<script>// javascript implementation of the above approach     var MAX_INDEX = 51;    var MAX_SUM = 2505;     // This dp array is used to store our values    // so that we don't have to calculate same    // values again and again    var dp = Array(MAX_INDEX).fill().map(()=>Array(MAX_SUM).fill().map(()=>Array(MAX_INDEX).fill(0)));;     function waysutil(index , sum , count, arr , K) {        // Base cases        // Index can't be less than 0        if (index < 0) {            return 0;        }         if (index == 0) {             // No element is picked hence            // average cannot be calculated            if (count == 0) {                return 0;            }            var remainder = sum % count;             // If remainder is non zero, we cannot            // divide the sum by count i.e. the average            // will not be an integer            if (remainder != 0) {                return 0;            }            var average = sum / count;             // If we find an average return 1            if (average == K) {                return 1;            }        }         // If we have already calculated this function        // simply return it instead of calculating it again        if (dp[index][sum][count] != -1) {            return dp[index][sum][count];        }         // If we don't pick the current element        // simple recur for index -1        var dontpick = waysutil(index - 1, sum, count, arr, K);         // If we pick the current element add it to        // our current sum and increment count by 1        var pick = waysutil(index - 1, sum + arr[index], count + 1, arr, K);        var total = pick + dontpick;         // Store the value for the current function        dp[index][sum][count] = total;        return total;    }     // Function to return the number of ways    function ways(N , K,  arr) {        var Arr = [];         // Push -1 at the beginning to        // make it 1-based indexing        Arr.push(-1);        for (i = 0; i < N; ++i) {            Arr.push(arr[i]);        }         // Initialize dp array by -1        for (i = 0; i < MAX_INDEX; i++) {            for (j = 0; j < MAX_SUM; j++) {                for (l = 0; l < MAX_INDEX; l++) {                    dp[i][j][l] = -1;                }            }        }         // Call recursive function        // waysutil to calculate total ways        var answer = waysutil(N, 0, 0, Arr, K);        return answer;    }     // Driver code             var arr = [ 3, 6, 2, 8, 7, 6, 5, 9 ];        var N = arr.length;        var K = 5;        document.write(ways(N, K, arr)); // This code contributed by gauravrajput1</script>
19
 
Time Complexity: O(MAX_INDEX * MAX_SUM * MAX_INDEX)Auxiliary Space: O(MAX_INDEX * MAX_SUM * MAX_INDEX) 
ankthon
princiraj1992
princi singh
GauravRajput1
singhh3010
Arrays
Dynamic Programming
Recursion
Arrays
Dynamic Programming
Recursion
Writing code in comment?
Please use ide.geeksforgeeks.org,
generate link and share the link here.
Maximum and minimum of an array using minimum number of comparisons
Top 50 Array Coding Problems for Interviews
Stack Data Structure (Introduction and Program)
Introduction to Arrays
Multidimensional Arrays in Java
0-1 Knapsack Problem | DP-10
Program for Fibonacci numbers
Longest Common Subsequence | DP-4
Bellman–Ford Algorithm | DP-23
Floyd Warshall Algorithm | DP-16