RVRSE - Editorial

PROBLEM LINK:

Problem Link

Author: Akshay Gahlot

Tester: Kritagya Agarwal

PREREQUISITES:

Probability, segment tree

PROBLEM:

Chef has an array A of N elements and he will perform the following operation K times:

  1. Select a subarray uniformly at random.
  2. Reverse the subarray.

Help him find the expected value at every index of the array A after he performs the operation K times.

Print the values as (P.Q^{-1}) mod 10^9+7, where Q^{−1} denotes the modular multiplicative inverse of Q (modulo 10^9+7).

QUICK EXPLANATION:

We can first solve the problem for a single swap, and then repeat the algorithms K times for final answer.

expVal[i]=P(A[i] is not swapped)*A[i]+P(i and j are swapped)*A[j])

EXPLANATION:

N^{2} solution for a single swap:

There are 2 cases that we need to handle:

  1. For 2 indices i and j, such that i <= j, the number of subarrays that if swapped, lead to i and j being swapped are min(i, N-j+1) (As the end points of subarray should be equidistant from both i and j.

  2. For an index i, the number of subarrays that if swapped, do not affect i are totalSubarrays - i * (N - i + 1). (i is in all the subarrays that have the left end \leq i and right end \geq i). Note that i swapped with itself is considered in the first case.

This gives an O(N^{2}) solution for a single swap.

To reduce this to O(Nlog(N)), we need to reduce the complexity of first step:

  1. For an index i, consider all indices j, such that j >= i. We can divide the expression min(i, N-j+1) into 2 parts:

    1. min(i, N-j+1) = i

    2. min(i, N-j+1) = N-j+1

    For the 1^{st} case, the range of j is i\leq N-j+1, ie, j \leq N-i+1. We can update the expected value of all these indices using range update on segment tree, ie update expected values from i+1 to N-i+1 by i*A[i] (ie, numSubarrays * A[i]).
    For the 2^{nd} case, we can update them when we iterate over to them.

  2. For an index j, consider all indices i, such that j > i. We need to update the expected value at index j, by all the elements in 2^{nd} part of the above step, ie all i such that i\gt N-j+1. So now we need to update the expected value at index j by (N-j+1)*A[i] (N-j+1 is the number of subarrays). This can be done by a range sum query over i. ie, sum=query(N-j+1 to j-1),

    expectedVal[j] += sum*(N-j+1).

The above steps are for i\leq j, so have to reverse the array and do the same. (taking care of equality)

The 2^{nd} part (index not swapped) is same as that of O(N^{2}) algorithm.

Now, we have to divide by total number of subarrays and print the solution in required format.

This is O(Nlog(N)) solution for 1 swap. Now we have to replace the array A by the expectedVal array, and repeat the algorithm K-1 times.

SOLUTIONS:

Setter's Solution
#include <bits/stdc++.h>
#define int long long
#define endl "\n"
#define MOD 1000000007

using namespace std;
		
int mpow(int a, int b){
    a %= MOD;
    if(!b)
        return 1;
    int temp = mpow(a, b / 2);
    temp = (temp * temp) % MOD;
    if(b % 2)
        return (a * temp) % MOD;
    return temp;
}

inline int mod_in(int n){
    return mpow(n, MOD - 2);
}

class segTree{
    public:
        int n;
        vector<int> st;
        vector<int> lazy;
        segTree(int x){
            n = x;
            st.resize(4 * n + 5, 0);
            lazy.resize(4 * n + 5, 0);
        }
    
        void push_down(int node, int l, int r) {
            if(lazy[node]){
            	lazy[node] %= MOD;
                st[node] += lazy[node] * (r - l + 1);
                st[node] %= MOD;
                if(r != l){
                    lazy[2 * node] += lazy[node];
                    lazy[2 * node + 1] += lazy[node];
                }
                lazy[node] = 0;
            }
        }

        void update(int lo, int hi, int val){
        	if (lo <= 0)
        		lo = 1;
        	if (hi > n)
        		hi = n;
            _up(1, 1, n, lo, hi, val);
        }
 
        int query(int lo, int hi){
            return _q(1, 1, n, lo, hi);
        }
 
        void _up(int node, int l, int r, int lo, int hi, int val){
            push_down(node, l, r);
            if (l > r || l > hi || lo > r)
                return;
            if (l >= lo && r <= hi) {
                lazy[node] += val;
                push_down(node, l, r);
                return;
            }
            int mid = (l + r) / 2;
            _up(2 * node, l, mid, lo, hi, val);    
            _up(2 * node + 1, mid + 1, r, lo, hi, val);
            st[node] = (st[2 * node] + st[2 * node + 1]) % MOD;
            return;
        }
 
        int _q(int node, int l, int r, int lo, int hi){
            push_down(node, l, r);
            if (l > r || l > hi || lo > r)
                return 0;
            if (l >= lo && r <= hi)
                return st[node];
            int mid = (l + r) / 2;
            int q1 = _q(2 * node, l, mid, lo, hi);    
            int q2 = _q(2 * node + 1, mid + 1, r, lo, hi);
            return (q1 + q2) % MOD;
        }
};

int32_t main(){

    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int n, k;
    cin >> n >> k;
    int a[n + 5];
    for (int i = 1 ; i <= n ; i++)
    	cin >> a[i];
    int tot_sub_arr = (n * (n + 1)) / 2;

    segTree exp_arr = segTree(n);
    for (int i = 1 ; i <= n ; i++) {
    	exp_arr.update(i, i, a[i]);
    }

    for (int i = 0 ; i < k ; i++) {
    	segTree st = segTree(n);
    	for (int i = 1 ; i <= n ; i++) {
    		int x = i + 1;
    		int y = n - i;
    		int val = exp_arr.query(i, i);
    		st.update(x, y, i * val);
    		x = n - i + 1;
    		y = i - 1;
    		st.update(i, i, (n - i + 1) * exp_arr.query(x, y));
    		y = i - 1;
    		x = n - i + 2;
    		st.update(x, y, (n - i + 1) * val);
    		x = i + 1;
    		y = n + 1 - i;
    		st.update(i, i, i * exp_arr.query(x, y));
    		int self = min(i, n + 1 - i);
    		int not_swapped = tot_sub_arr - i * (n + 1 - i);

    		st.update(i, i, exp_arr.query(i, i) * (self + not_swapped));
    	}
    	swap(st, exp_arr);
    }

    int denom = mod_in(mpow(tot_sub_arr, k));

    for (int i = 1 ; i <= n ; i++) {
        cout << (exp_arr.query(i, i) * denom ) % MOD << " ";
    }
}
Tester's Solution
#include<bits/stdc++.h>
#define ll long long int
#define mod 1000000007
using namespace std;
ll a[1001];
 
ll power(ll a, ll n)
{
    if(n == 0) return 1;
 
    ll ans = 1;
    ll val = a;
    while(n)
    {
        if(n&1)
        {
            ans *= val;
            ans %= mod;
        }
        val *= val;
        val %= mod;
        n /= 2;
    }
 
    return ans;
}
 
void apply(ll n)
{
    ll pref[n+1];
    ll suff[n+1];
 
    ll total = n*(n+1)/2;
    total %= mod;
 
    ll nt[n+1];
 
    ll b[n+1];
    ll sum[n+1] = {0};
 
    for(int i = 1 ; i <= n ; i++)
    {
        sum[i] = a[i] + sum[i-1];
        sum[i] %=  mod;
    }
 
    for(int i = 1 ; i <= n ; i++)
    {
        b[i] = a[n - i + 1];
    }
 
    pref[0] = suff[0] = 0;
 
    for(int i = 1 ; i <= n ; i++)
    {
        pref[i] = (i*a[i]) % mod;
        suff[i] = (i*b[i]) % mod;
    }
 
    for(int i = 1 ; i <= n ; i++)
    {
        pref[i] += pref[i-1];
        pref[i] %= mod;
 
        suff[i] += suff[i-1];
        suff[i] %= mod;
    }
 
    ll answer[n+1] = {0};
    int l = 2;
    int r = n - 1;
    for(int i = 1 ; i <= n/2 ; i++)
    {
        ll ans = (pref[i] + suff[i]) % mod;
 
        ll val1 = 0;
 
        if(l <= r)
        {
            val1 = i;
            val1 *= (sum[r] - sum[l-1] + mod)%mod;
            val1 %= mod;
        }
        ans += val1;
        ans %= mod;
 
        l++;
        r--;
 
        answer[i] = ans;
        answer[n - i + 1] = ans;
 
        ll temp = ((n-i)*i)%mod;
        temp += i;
        temp %= mod;
 
        temp = total - temp + mod;
        temp %= mod;
        nt[i] = temp;
        nt[n - i + 1] = temp;
    }
 
    if(n&1)
    {
        int index = n/2 + 1;
        answer[index] = pref[index] + suff[index - 1];
        answer[index] %= mod;
 
        ll ty = n/2;
        nt[index] = (ty*(ty+1))% mod;
        nt[index] += index;
        nt[index] %= mod;
 
        nt[index] = total - nt[index] + mod;
        nt[index] %= mod;
    }
 
 
    for(int i = 1 ; i <= n ; i++)
    {
        answer[i] += (a[i]*nt[i])%mod;
        answer[i] %= mod;
 
        answer[i] *= power(total, mod-2);
        answer[i] %= mod;
    }
    
    for(int i = 1 ; i <= n ; i++)
    {
        a[i] = answer[i];
    }
}
 
int main(){
    ll n;
    cin >> n;
 
    ll k;
    cin >> k;
 
    for(int i = 1 ; i <= n ; i++){
        cin >> a[i];
    }
 
    for(int i = 1 ; i <= k ; i++)
    { 
        apply(n);
    }
 
    for(int i = 1 ; i <= n ; i++)
    {
        cout << a[i] << " ";
    }
    cout << endl;
}
2 Likes

You don’t need a segment tree. Similar to your approach, I was trying to find the expected value of the i^{th} index. Suppose we want to bring some j on i. Also, suppose j < i, then we have the sum: \sum\limits_{j=1}^{i-1} min(j, n-i+1)*arr[j]. If i \leq n-i+1, then we have \sum\limits_{j=1}^{i-1} j*arr[j]. If i > n-i+1, we then have: \sum\limits_{j=1}^{n-i}j*arr[j] + \sum\limits_{j=n-i+1}^{i-1} (n-i+1)*arr[j]. All these can be handled using prefix sums. Similarly, we can find the expressions for j > i. Unfortunately, i couldn’t get AC :slightly_frowning_face:.
Overall : O(N*K).

Tester’s solution uses only prefix-suffix sums as well. However, I felt segment tree is more convenient to use.