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:
- Select a subarray uniformly at random.
- 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:
-
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.
-
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:
-
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:
-
min(i, N-j+1) = i
-
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. -
-
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;
}