PROBLEM LINK:
Author: Ritesh Gupta
Tester: Jatin Yadav
Editorialist: Ritesh Gupta
DIFFICULTY:
MEDIUM
PREREQUISITES:
PREFIX SUM, FENWICK TREE, SEGMENT TREE
PROBLEM:
You are given a circular array A(|A_i| \le 10^9) of N(1 \le N \le 2*10^5) integers. A_{i,j} defined as an array formed using elements of A from i^{th} index element to j^{th} index element(both inclusive) in a clockwise direction.
A function f(A_{i,j}) defined as the count of a subarray of array A_{i,j} whose sum is less than K(|K| \le 10^{15}). You have to find the sum of f(A_{i,j}) for all valid i and j under modulo 10^9 + 7.
QUICK EXPLANATION:
- The solution is explained here in 1- based indexing.
- Let, first solve the problem for the non-circular array. For that find the number of subarrays that end at i^{th} index and have the sum less than K. It can be done by finding the count of indexes of the prefixes that have ended before it and have the sum such that if we subtract that from the i^{th} index prefix then it evaluates less than K.
- Suppose there are D_i prefixes d_{i,1}, d_{i,2}, ..., d_{i,D} are of the type above mentioned for index i and the contribution of all the subarray who ends at index i is then \sum_{j = 1}^{D_i} d_{i,j} \space * (N-i +1).
- Then the answer to for non-circular array can be written as \sum_{i = 1}^{N} \sum_{j = 1}^{D_i} d_{i,j} \space * (N-i +1).
EXPLANATION:
Let’s solve the problem in two parts:
-
Calculating the contribution of subarray which constructed by eliminating the two prefixes of the given array.
As there are (n*(n+1))/2 subarray of this type so we could use the same technique mentioned in the quick explanation. Let’s suppose two prefixes ending at i^{th} and j^{th} has sum s_i and s_j such that j < i. Then how many subarrays (consider all possible subarray including circular one) consist of this as it’s a subarray. The answer to this is given by the following sum:
- 1 + 2 + ... + (N-i+j+1)
- ((N-i+j+1)*(N-i+j+2))/2
- ((N-i+1)*(N-i+2) + (2*N-2*i+3)*j + j*j)
If there are cnt possible values for j then the answer is given by:
- (cnt*(N-i+1)*(N-i+2) + (2*N-2*i+3)*\sum_{j \subset k} j + \sum_{j \subset k} j*j)
We need to do this process for each index and this can be efficiently done by Fenweek Tree or segment Tree.
-
Calculating the contribution of subarray which constructed by the concatenation of non-empty suffix and non-empty prefix of the given array.
To do that we calculate the valid prefixes for each valid i such that if we consider the suffix from (i+1) to N such that sum is less than K. We can do this in the same way as described above. For better understanding refer to the setter solution.
COMPLEXITY:
TIME: O(NlogN)
SPACE: O(NlogN)
SOLUTIONS:
Setter's Solution
#include <bits/stdc++.h>
#define int long long
#define mod 1000000007
using namespace std;
const int N = 200010;
const int cat = 500000004;
int bit1[N][3],a[N],c[N];
int m;
void update(int idx, int val)
{
while(idx<=m)
{
bit1[idx][0]+=val;
bit1[idx][1]+=(val*val);
bit1[idx][2]++;
idx+=idx&-idx;
}
}
pair<int,pair<int,int> > pref(int idx)
{
pair<int,pair<int,int> > ans = {0,{0,0}};
while(idx>0)
{
ans.first+=bit1[idx][0];
ans.second.first+=bit1[idx][1];
ans.second.second+=bit1[idx][2];
idx-=idx&-idx;
}
return ans;
}
pair<int,pair<int,int> > rsum(int l, int r)
{
auto c1 = pref(r);
auto c2 = pref(l-1);
c1.first -= c2.first;
c1.second.first -= c2.second.first;
c1.second.second -= c2.second.second;
return c1;
}
int32_t main()
{
int t;
cin >> t;
while(t--)
{
int n,k;
cin >> n >> k;
for(int i=0;i<n;i++)
cin >> a[i];
set <int> s;
s.insert(0);
int sum = 0;
for(int i=0;i<n;i++)
sum += a[i], s.insert(sum);
int i = 1;
map <int,int> pos;
for(int j:s)
pos[j] = i++;
m = pos.size();
for(int i=0;i<=m;i++)
for(int j=0;j<3;j++)
bit1[i][j] = 0;
update(pos[0],0);
sum = 0;
int ans = 0;
for(int i=0;i<n;i++)
{
sum += a[i];
auto cnt = pos.lower_bound(sum-k+1);
if(cnt == pos.end())
{
update(pos[sum],i+1);
continue;
}
auto temp = rsum(cnt->second,pos.size());
int pat = temp.first + temp.second.first + (2*(n-i)*temp.first)%mod + (temp.second.second*(n-i)*(n-i+1))%mod;
pat = ((pat%mod) * (cat%mod))%mod;
ans = (ans + pat)%mod;
// ans += (temp.first + temp.second.first + 2*(n-i)*temp.first + temp.second.second*(n-i)*(n-i+1))/2;
update(pos[sum],i+1);
}
for(int i=0;i<=m;i++)
for(int j=0;j<3;j++)
bit1[i][j] = 0;
c[n-1] = a[n-1];
for(int i=n-2;i>=0;i--)
c[i] = c[i+1] + a[i];
sum = 0;
for(int i=0;i<n-1;i++)
{
sum += a[i];
update(pos[sum],i);
auto cnt = pos.upper_bound(k-c[i+1]-1);
if(cnt == pos.begin())
continue;
cnt--;
auto temp = pref(cnt->second);
int pat = ((i+1)*(i+2)*temp.second.second)%mod - ((2*i+3)*temp.first)%mod + temp.second.first%mod + mod;
pat = ((pat%mod) * (cat%mod))%mod;
ans = (ans + pat)%mod;
// ans += ((i+1)*(i+2)*temp.second.second - (2*i+3)*temp.first + temp.second.first)/2;
}
cout << ans%mod << endl;
}
}
Tester's Solution
#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define sd(x) scanf("%d", &(x))
#define pii pair<int, int>
#define F first
#define S second
#define all(c) ((c).begin()), ((c).end())
#define sz(x) ((int)(x).size())
#define ld long double
template<class T,class U>
ostream& operator<<(ostream& os,const pair<T,U>& p){
os<<"("<<p.first<<", "<<p.second<<")";
return os;
}
template<class T>
ostream& operator <<(ostream& os,const vector<T>& v){
os<<"{";
for(int i = 0;i < (int)v.size(); i++){
if(i)os<<", ";
os<<v[i];
}
os<<"}";
return os;
}
#ifdef LOCAL
#define cerr cout
#else
#endif
#define TRACE
#ifdef TRACE
#define trace(...) __f(#__VA_ARGS__, __VA_ARGS__)
template <typename Arg1>
void __f(const char* name, Arg1&& arg1){
cerr << name << " : " << arg1 << std::endl;
}
template <typename Arg1, typename... Args>
void __f(const char* names, Arg1&& arg1, Args&&... args){
const char* comma = strchr(names + 1, ',');cerr.write(names, comma - names) << " : " << arg1<<" | ";__f(comma+1, args...);
}
#else
#define trace(...)
#endif
const int M = 4e5 + 10;
struct PIII{
int s0; ll s1, s2;
PIII operator += (const PIII & p){
s0 += p.s0;
s1 += p.s1;
s2 += p.s2;
};
PIII operator + (const PIII & p) const{
return {s0 + p.s0, s1 + p.s1, s2 + p.s2};
};
};
template<class T>
struct Data{
int child[2], height;
T bits[2];
PIII val;
Data(){
child[0] = bits[0] = child[1] = bits[1] = height = 0;
val = {0, 0LL, 0LL};
}
};
Data<ll> D[M]; // ll
int cnt = 0;
void init(){
for(int i = 0; i <= cnt; i++) D[i] = Data<ll>();
cnt = 0;
}
inline int clz(ll N) { // ll
return N ? 63 - __builtin_clzll(N) : -1;
}
const ll OFFSET = 1LL << 50;
template<class T, int num_bits>
struct compressed_trie{
int root;
PIII sm;
compressed_trie(){
root = ++cnt;
D[root].height = num_bits;
sm = {0, 0LL, 0LL};
}
void add(T x, PIII v){
x += OFFSET;
sm += v;
int node = root;
int pos = num_bits - 1;
for(; ; ){
Data<T> & d = D[node];
d.val += v;
if(pos < 0) return;
int dir = x >> pos & 1;
if(!d.child[dir]){
d.child[dir] = ++cnt;
D[cnt].val += v;
d.bits[dir] = x;
return;
}
T y = d.bits[dir];
int nxt = d.child[dir];
int nxtH = D[nxt].height;
y <<= nxtH;
int diff = clz(x^y);
if(diff < nxtH){
node = nxt;
pos = nxtH - 1;
x &= ((T(1)<<nxtH) - 1);
continue;
}
int yside = y >> diff & 1;
int xside = x >> diff & 1;
d.child[dir] = ++cnt;
d.bits[dir] = x >> (diff + 1);
D[cnt].val = D[nxt].val + v;
D[cnt].height = diff + 1;
D[cnt].child[yside] = nxt;
D[cnt].bits[yside] = (y >> nxtH) & ( (T(1) << (diff + 1 - nxtH)) - 1);
D[cnt].bits[xside] = x & ((T(1) << (diff + 1)) - 1);
D[cnt].child[xside] = cnt + 1; ++cnt;
D[cnt].val = v;
return;
}
}
PIII order_of_key(T x){
x += OFFSET;
int node = root;
int pos = num_bits - 1;
PIII ret = {0, 0LL, 0LL};
for(; pos >= 0; ){
int dir = x >> pos & 1;
Data<T> & d = D[node];
if(dir) ret += D[d.child[0]].val;
if(!d.child[dir]){
return ret;
}
T y = d.bits[dir];
int nxt = d.child[dir];
int nxtH = D[nxt].height;
y <<= nxtH;
int diff = clz(x^y);
if(diff < nxtH){
node = nxt;
pos = nxtH - 1;
x &= ((T(1)<<nxtH) - 1);
continue;
}
if(x >> diff & 1) ret += D[nxt].val;
return ret;
}
return ret;
}
void insert(T x){
add(x, 1);
}
};
const int mod = 1e9 + 7;
int main(){
int t; sd(t);
while(t--){
init();
int n; ll k; sd(n); scanf("%lld", &k);
compressed_trie<ll, 55> ct;
vector<ll> pref(n + 1);
ll ALL = 0, ans = 0;
for(int i = 1; i <= n; i++){
int x; sd(x);
ALL += x;
pref[i] = pref[i - 1] + x;
}
for(int j = 1; j <= n; j++){
ll thresh = pref[j] - k + 1;
PIII SM = ct.sm;
PIII V = ct.order_of_key(thresh);
ll s0 = SM.s0 - V.s0;
ll s1 = (SM.s1 - V.s1) % mod;
ll s2 = (SM.s2 - V.s2) % mod;
if(0 >= thresh) s0++;
ll t = n + 2 - j;
ans += (t * (t - 1) % mod) * s0 + (2 * t - 1) * s1 + s2;
{
ll thresh = k + pref[j - 1] - ALL;
PIII V = ct.order_of_key(thresh);
ll t = j + 1;
ans += (t * (t - 1) % mod) * V.s0 - (2 * t - 1) * V.s1 + V.s2;
}
ans %= mod;
ct.add(pref[j], {1, j, j * (ll) j});
}
ans = (ans % mod + mod) % mod;
printf("%lld\n", (ans * ((mod + 1) / 2)) % mod);
}
}