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.

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.