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

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 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: 19

Input: 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 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++

filter_none

edit
close

play_arrow

link
brightness_4
code

#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 again
int 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 ways
int 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 code
int 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;
}

chevron_right


Java

filter_none

edit
close

play_arrow

link
brightness_4
code

// Java implementation of the above approach
import 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 */

chevron_right


Python3

filter_none

edit
close

play_arrow

link
brightness_4
code

# Python implementation of above approach
import numpy as np
  
MAX_INDEX = 51
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 
  
# Initialize dp array by -1 
dp = 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 ways 
def 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 code 
if __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

chevron_right


C#

filter_none

edit
close

play_arrow

link
brightness_4
code

// C# implementation of the above approach
using 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

chevron_right


Output:

19


My Personal Notes arrow_drop_up

Check out this Author's contributed articles.

If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

Please Improve this article if you find anything incorrect by clicking on the "Improve Article" button below.