Search code examples
c++arraysalgorithmdynamic-programmingcounting

Efficient algorithm to compute combinations with repetitions of an array adding up to given sum


So in a personal C++ project I am faced with a problem. I am rephrasing it as follows :

Given an array of n elements (e.g. [1, 3, 5], with n = 3 elements) where the number at the ith position denotes how many possible values the number at ith index can take (e.g here the first element can take 1 value which is 0; the second element can take 3 values from among 0,1,2; the third element can take 5 values from among 0,1,2,3,4).

I need to list all possible such arrays of length n that sum up to less than or equal to a given number k. Here is an example :

Input 1:

input array = [2,2]; k = 2

Output 1:

[0,0], [0,1], [1,0], [1,1]

Also, for instance :

Input 2:

input array = [2,2]; k = 1

Output 2:

[0,0], [0,1], [1,0]

The issue :

I have coded up a simple recursive and a simple iterative solution, which enumerates all arrays and only keeps those which have sum less than k. The problem with these is that for the case where n is large and k = 1, my code takes very long to run, since it enumerates all the cases and keeps a few.

I cannot see any overlapping sub-problems so I feel DP and memoization isn't applicable. How can I write the required C++ code for this that works?

Here is my code for the iterative version :

// enumerates all arrays which sum up to k

vector<vector<int> > count_all_arrays(vector<int> input_array, int k){

    vector<vector<int> > arr;
    int n = (int)input_array.size();

    // make auxilliary array with elements

    for(int i = 0; i < n; i++){
        vector<int> temp(input_array[i]);
        std::iota(temp.begin(), temp.end(), 0);
        arr.push_back(temp);
    }

    // computes combinations

    vector<int> temp(n);
    vector<vector<int> > answers;
    vector<int> indices(n, 0);
    int next;

    while(1){ 
        temp.clear();
        for (int i = 0; i < n; i++) 
            temp.push_back(arr[i][indices[i]]);  
        long long int total = accumulate(temp.begin(), temp.end(), 0);
        if(total <= k)
            answers.push_back(temp);
        next = n - 1; 
        while (next >= 0 &&  
              (indices[next] + 1 >= (int)arr[next].size())) 
            next--; 
        if (next < 0) 
            break; 
        indices[next]++; 
        for (int i = next + 1; i < n; i++) 
            indices[i] = 0; 
    }
    return answers;
}

Solution

  • It's a pretty simple recursive task:

    #include <bits/stdc++.h>    
    using namespace std;
    
    int arr[] = {2, 2};
    int n = 2;
    int k = 2;
    
    void gen(int pos, int sum, string s){
        if(pos == n){
            cout<<"["<<s<<" ]"<<endl;
            return;
        }
        for(int i = 0; i < arr[pos]; i++){
            if(sum + i > k) return;
            gen(pos + 1, sum + i, s + " " + to_string(i));
        }
    }
    
    int main(){
        gen(0, 0, "");
        return 0;
    }
    

    Just generate all possibilities for each slot of the array and for each choice, take the sum to the evaluation of the next slot.

    When n is large and k = 1, it's natural that it takes O(n), since you will have:

    [0, 0, 0, ..., 0, 0, 1]
    [0, 0, 0, ..., 0, 1, 0]
    [0, 0, 0, ..., 1, 0, 0]
              ...
    [0, 0, 1, ..., 0, 0, 0]
    [0, 1, 0, ..., 0, 0, 0]
    [1, 0, 0, ..., 0, 0, 0]