Number of ways to write n as sum of k numbers with restrictions on each part

Question

Title says it all.

I need to split n as sum of k parts where each part k_i should be in the range of 1 <= k_i <= r_i for given array r.

for example -

n = 4, k = 3 and r = [2, 2, 1]
ans = 2
#[2, 1, 1], [1, 2, 1]

Order matters. (2, 1, 1) and (1, 2, 1) are different.

I taught of solving it using stars and bars method, but be because of upper bound r_i i dont know to to approach it.

i implemented a direct recursion function and it works fine for small values only.

Constraints of original problem are

1 <= n <= 107

1 <= k <= 105

1 <= ri <= 51

All calculations will be done under prime Modulo.

i found a similar problem here but i don't know how to implement in program. HERE

My brute-force recursive function -

#define MAX 1000
const int md = 1e9 + 7;

vector <int> k;
vector <map<int, int>> mapper;

vector <int> hold;

int solve(int sum, int cur){

    if(cur == (k.size() - 1) &&  sum >= 1 && sum <= k[cur]) return 1;
    if(cur == (k.size() - 1) &&  (sum < 1 || sum > k[cur])) return 0;

    if(mapper[cur].find(sum) != mapper[cur].end())
        return mapper[cur][sum];

    int ans = 0;
    int start = 1;

    for(int i=start; i<=k[cur]; ++i){


        int remain = sum - i;
        int seg = (k.size() - cur) - 1;
        if(remain < seg) break;

        int res = solve(sum - i, cur + 1);
        ans = (1LL * ans + res) % md;
    }

    mapper[cur][sum] = ans;
    return ans;
}


int main(){

    for(int i=0; i<MAX; ++i) k.push_back(51);  // restriction for each part default 51
    mapper.resize(MAX);

    cout << solve(MAX + MAX, 0) << endl;
}

Instead of using a map for storing result of computation i used a two dimensional array and it gave very good performance boost but i cannot use it because of large n and k values.

How could i improve my recursive function or what are other ways of solving this problem.

Michael

Can you please share link of the original problem, if possible? The question becomes much simpler if upper bound is same for every partition but apparently it isn't from the example you provided. – Badminton

Is k always equal to the size of r? – Guyguyana

@Guyguyana Yes always. – Michael

Answer 1

That's interesting problem.

First lets say r_i = r_i - 1, n = n - k, numbers in [0, r_i] just for convenience. Now it's possible to add some fictitious numbers to make m the power of 2 without changing answer.

Now let's represent each interval of [0, r_i] as polynomial 1 * x ^ 0 + 1 * x ^ 1 + ... + 1 * x & r_i. Now if we multiply all these polynomials, coefficient at x ^ n will be answer.

Here is structure called Number Theoretic Transform (NTT) which allows to multiply two polynomials modulo p in O(size * log(size)).

If you will just multiply it using NTT, code will work in something like O(n * k * log (k * max(r))). It's very slow.

But now our fictive numbers help. Let's use divide and conquer technics. We'll make O(log m) steps, on each step multiply 2 * i-th and 2 * i + 1-th polynomials. In the next step we'll multiply resulting polynomials of this step.

Each step works in O(k * log(k)) and there is O(log(k)) steps, so algorhitm works in O(k * log^2 (k)). It's fast asymptotically, but I'm not sure if it fits TL for this problem. I think it will work about 20 seconds on max test.

Underproduction