EXPPROD - Editorial

PROBLEM LINK:

Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4

Author: raysh07
Tester: sushil2006
Editorialist: iceknight1093

DIFFICULTY:

TBD

PREREQUISITES:

Dynamic programming, Mo’s algorithm

PROBLEM:

For an array A, f(A, K) is defined as follows:

  • K times, randomly increase an element of A.
  • f(A, K) is defined to be the expected product of A after all K increments.

You’re given an array A.
Answer Q queries; each of which gives you L, R, K and asks for f(A[L, R], K).

EXPLANATION:

First, we need to understand how to compute f(A, K) for a fixed array A.

Let x_i denote the number of times the value at index i is increased.
So, any final state of the array can be described by the tuple (x_1, \ldots, x_n), which will be denoted (x) for convenience.

Our aim is to compute the expectation of

\prod_{i=1}^N \left(A_i + x_i \right)

across all possible valid tuples (x).

It can be proved that this expected value is

\frac{K!}{N^K} \cdot \left( \sum_{S \subseteq \{1, \ldots, N\}} \left(\prod_{j\not\in S} A_j\right) \frac{N^{K-|S|}}{(K-|S|)!}\right)

which, by grouping sets of the same size (so |S| above) can be further reduced to

\sum_{y=0}^{\min(N, K)} c_m \cdot \frac{K!}{(K-m)!} \cdot \frac{1}{N^m}

where

c_m = \sum_{|S| = m} \left(\prod_{j\not\in S} A_j \right)

Observe that if the c_m values are known, the answer can be computed in linear time.

Further, note that c_m is exactly the sum of products of all size-(N-m) subsets of A, and so can be computed quite easily (in quadratic time) using a knapsack.

This allows us to solve for a single query in quadratic time.


Now, we need to optimize this to solving multiple queries.

The bottleneck here is computing the coefficients c_m.
Note that if we have the coefficients for the range [L, R] already, extending them to [L-1, R] or [L, R+1] takes \mathcal{O}(N) time, since we only need to process the insertion of a single element (either A_{L-1} or A_R, respectively).
Further, since this is a knapsack-style DP, it’s also possible to remove a single element in linear time as well!
(See point 5 of this blog if you’re unfamiliar with the idea).

Thus, from having the coefficients for [L, R] we can move any of L or R by \pm 1 in linear time.

This allows us to simply apply Mo’s algorithm, leading to a complexity of \mathcal{O}((N+Q)\sqrt{Q}\cdot N) which is fast enough in practice.

TIME COMPLEXITY:

\mathcal{O}((N+Q)\sqrt{Q}\cdot N) per testcase.

CODE:

Author's code (C++)
#include <bits/stdc++.h>
using namespace std;
#define int long long
#define INF (int)1e18

mt19937_64 RNG(chrono::steady_clock::now().time_since_epoch().count());

const int mod = 1e9 + 7;

struct mint{
    int x;

    mint (){ x = 0;}
    mint (int32_t xx){ x = xx % mod; if (x < 0) x += mod;}
    mint (long long xx){ x = xx % mod; if (x < 0) x += mod;}

    int val(){
        return x;
    }
    mint &operator++(){
        x++;
        if (x == mod) x = 0;
        return *this;
    }
    mint &operator--(){
        if (x == 0) x = mod;
        x--;
        return *this;
    }
    mint operator++(int32_t){
        mint result = *this;
        ++*this;
        return result;
    }
    
    mint operator--(int32_t){
        mint result = *this;
        --*this;
        return result;
    }
    mint& operator+=(const mint &b){
        x += b.x;
        if (x >= mod) x -= mod;
        return *this;
    }
    mint& operator-=(const mint &b){
        x -= b.x;
        if (x < 0) x += mod;
        return *this;
    }
    mint& operator*=(const mint &b){
        long long z = x;
        z *= b.x;
        z %= mod;
        x = (int)z;
        return *this;
    }
    mint operator+() const {
        return *this;
    }
    mint operator-() const {
        return mint() - *this;
    }
    mint operator/=(const mint &b){
        return *this = *this * b.inv();
    }
    mint power(long long n) const {
        mint ok = *this, r = 1;
        while (n){
            if (n & 1){
                r *= ok;
            }
            ok *= ok;
            n >>= 1;
        }
        return r;
    }
    mint inv() const {
        return power(mod - 2);
    }
    friend mint operator+(const mint& a, const mint& b){ return mint(a) += b;}
    friend mint operator-(const mint& a, const mint& b){ return mint(a) -= b;}
    friend mint operator*(const mint& a, const mint& b){ return mint(a) *= b;}
    friend mint operator/(const mint& a, const mint& b){ return mint(a) /= b;}
    friend bool operator==(const mint& a, const mint& b){ return a.x == b.x;}
    friend bool operator!=(const mint& a, const mint& b){ return a.x != b.x;}
    mint power(mint a, long long n){
        return a.power(n);
    }
    friend ostream &operator<<(ostream &os, const mint &m) {
        os << m.x;
        return os;
    }
    explicit operator bool() const {
        return x != 0;
    }
};

// Remember to check MOD

struct factorials{
    int n;
    vector <mint> ff, iff;
    
    factorials(int nn){
        n = nn;
        ff.resize(n + 1);
        iff.resize(n + 1);
        
        ff[0] = 1;
        for (int i = 1; i <= n; i++){
            ff[i] = ff[i - 1] * i;
        }
        
        iff[n] = ff[n].inv();
        for (int i = n - 1; i >= 0; i--){
            iff[i] = iff[i + 1] * (i + 1);
        }
    }
    
    mint C(int n, int r){
        if (n == r) return mint(1);
        if (n < 0 || r < 0 || r > n) return mint(0);
        return ff[n] * iff[r] * iff[n - r];
    }
    
    mint P(int n, int r){
        if (n < 0 || r < 0 || r > n) return mint(0);
        return ff[n] * iff[n - r];
    }
    
    mint solutions(int n, int r){
        // Solutions to x1 + x2 + ... + xn = r, xi >= 0 
        return C(n + r - 1, n - 1);
    }
    
    mint catalan(int n){
        return ff[2 * n] * iff[n] * iff[n + 1];
    }
};

const int PRECOMP = 5000;
factorials F(PRECOMP);

// REMEMBER To check MOD and PRECOMP

void Solve() 
{
    int n, q; cin >> n >> q;
    
    vector <int> a(n + 1);
    for (int i = 1; i <= n; i++){
        cin >> a[i];
    }
    
    vector <mint> dp(n + 1);
    dp[0] = 1;
    
    vector <array<int, 4>> qq(q);
    vector <mint> res(q);
    for (int i = 0; i < q; i++){
        cin >> qq[i][0] >> qq[i][1] >> qq[i][2];
        qq[i][3] = i;
    }
    
    const int B = sqrt(n);
    
    sort(qq.begin(), qq.end(), [&](array <int, 4> x, array <int, 4> y){
        if (x[0] / B != y[0] / B){
            return x[0] < y[0];
        }
        return x[1] < y[1];
    }); 
    
    auto add = [&](int x){
        mint y = a[x];  
        
        for (int i = n; i >= 1; i--){
            dp[i] += dp[i - 1] * y;
        }
    };
    
    auto era = [&](int x){
        mint y = a[x];
        
        for (int i = 1; i <= n; i++){
            dp[i] -= dp[i - 1] * y;
        }
    };
    
    int L = 1, R = 1;
    add(1);
    
    vector <mint> inv(n + 1);
    for (int i = 1; i <= n; i++){
        inv[i] = mint(1) / i;
    }
    
    for (auto [l, r, k, i] : qq){
        while (L > l){
            L--;
            add(L);
        }
        while (R < r){
            R++;
            add(R);
        }
        while (L < l){
            era(L);
            L++;
        }
        while (R > r){
            era(R);
            R--;
        }
        
        mint ans = 0;
        mint choose = 1;
        mint ok = 1;
        mint to_inv = mint(1) / (r - l + 1);
        for (int i = 0; i <= r - l + 1; i++){
            mint val = dp[(r - l + 1) - i];
            
            if (i > 0){
                choose *= (k - i + 1);
                choose *= inv[i];
                ok *= to_inv;
            }
            
            mint term = val * choose * ok * F.ff[i];
            ans += term;
        }
        
        res[i] = ans;
    }
    
    for (int i = 0; i < q; i++){
        cout << res[i] << "\n";
    }
}

int32_t main() 
{
    auto begin = std::chrono::high_resolution_clock::now();
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    int t = 1;
    // freopen("in",  "r", stdin);
    // freopen("out", "w", stdout);
    
    cin >> t;
    for(int i = 1; i <= t; i++) 
    {
        //cout << "Case #" << i << ": ";
        Solve();
    }
    auto end = std::chrono::high_resolution_clock::now();
    auto elapsed = std::chrono::duration_cast<std::chrono::nanoseconds>(end - begin);
    cerr << "Time measured: " << elapsed.count() * 1e-9 << " seconds.\n"; 
    return 0;
}

We Can use a Segement Tree Built on Convolution.
Look at this submission Link