Codeforces 992E - Nastya and King-Shamans

Problem: Codeforces 992E Nastya and King-Shamans
Given n intergers $ a_{1},a_{2},…a_{n} $ and q operations $(p_{i}, x_{i})$, $1 \le n,q \le 2 \cdot 10^{5} $. In each operation, change $a_{p_{i}}$ to $x_{i}$, and after that give any index k such that $ a_{k} = \sum_{i=1}^{k-1} a_{i} $ if there exists one.

Segment tree solution

For each query, we search the answer from left to right. We maintain a cur_sum, which is the sum of $a_{i}$ we have passed, meaning that we haven’t found a answer till now. So we go on to find it in the right part. We need to find an leftmost element that is at least cur_sum. If we find one, then we are done. If we don’t find it, we add that element to cur_sum. As it is the leftmost one, we wouldn’t omit any answer. And since the cur_sum will be at least twice as big as it was, the count repeated search for one operation will be at most log(MAX).
We can use segment tree to maintain the sum and max value to do the queries.
My code:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
#include <iostream>
#include <algorithm>
#include <string.h>
#include <math.h>

using namespace std;

#define pb push_back
#define mp make_pair
#define ll long long
#define ull unsigned ll
#define db double
#define INF 0x3f3f3f3f
#define MOD 1000000007
#define PII pair<int, int>

const int N=200010;
ll s[N*4],mx[N*4],a[N];
int n,q,p,x;

void build(int x,int l,int r) {
if (l==r) {
mx[x]=s[x]=a[l];
return;
}
int m=(l+r)>>1;
build(2*x,l,m);
build(2*x+1,m+1,r);
mx[x]=max(mx[2*x],mx[2*x+1]);
s[x]=s[2*x]+s[2*x+1];
}

void update(int x,int l,int r,int pos,ll val) {
if (l==r) {
mx[x]=s[x]=val;
return;
}
int m=(l+r)>>1;
if (pos<=m) update(2*x,l,m,pos,val);
else update(2*x+1,m+1,r,pos,val);
mx[x]=max(mx[2*x],mx[2*x+1]);
s[x]=s[2*x]+s[2*x+1];
}

ll get_sum(int x,int l,int r,int lq,int rq) {
if (lq<=l&&r<=rq) return s[x];
if (lq>r||rq<l) return 0;
int m=(l+r)>>1;
return get_sum(2*x,l,m,lq,rq)+get_sum(2*x+1,m+1,r,lq,rq);
}

int down(int x,int l,int r,ll val) {
if (l==r) {
if (mx[x]>=val) return l;
else return n+1;
}
int m=(l+r)>>1;
if (mx[2*x]>=val) return down(2*x,l,m,val);
return down(2*x+1,m+1,r,val);
}

int get_left(int x,int l,int r,int lq,int rq,ll val) {
if (l>rq||lq>r) return n+1;
if (lq<=l&&r<=rq) {
if (mx[x]<val) return n+1;
return down(x,l,r,val);
}
int m=(l+r)>>1;
int res=get_left(2*x,l,m,lq,rq,val);
if (res!=n+1) return res;
return get_left(2*x+1,m+1,r,lq,rq,val);
}

int solve() {
int idx=0;
ll sum=0;
while (idx<n+1) {
int new_idx=get_left(1,1,n,idx+1,n,sum);
if (new_idx==n+1) return -1;
ll pre_sum=get_sum(1,1,n,1,new_idx-1);
ll ele=get_sum(1,1,n,new_idx,new_idx);
if (pre_sum==ele) return new_idx;
idx=new_idx;
sum=pre_sum+ele;
}
return -1;
}

int main()
{
scanf("%d%d",&n,&q);
for (int i=1;i<=n;i++) scanf("%lld",a+i);
build(1,1,n);
for (int i=1;i<=q;i++) {
scanf("%d%d",&p,&x);
update(1,1,n,p,x);
printf("%d\n",solve());
}
}

Square decomposition

Let $p_{i}$ be the prefix sum of a, and $f_{i} = p_{i} - 2 \cdot a_{i}$. Then we only need to find an i such that $f_{i} = 0$. And each time we change $a_{i}$ to val, let $\delta = val - a_{i}$, $f_{i}$ should decrease by $\delta$ whereas $f_{k}(k>i)$ should increase by $\delta$.
Next we divide the sequence of groups of size M. For each change, we may change at most N\M groups as a whole and deal with all the elements in one group. And for query, we keep all elements in a group in order so that we can use binary search. To do with the single group which is not entirely coverd by the range, we can do something like merging sorted vectors.
Though the solution is still too slow to pass the tests, I think square decomposition is powerful in data structrue problems.
My code:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
#include <iostream>
#include <algorithm>
#include <string.h>
#include <math.h>
#include <vector>

using namespace std;

#define pb push_back
#define mp make_pair
#define ll long long
#define ull unsigned ll
#define db double
#define INF 0x3f3f3f3f
#define MOD 1000000007
#define PII pair<int, int>

const int N=200010;
int M=500;
int n,p,q,x,a[N],s[N];
struct node {
ll val;
int idx;
bool operator<(const node& rhs) const {
return val<rhs.val;
}
bool operator==(const node& rhs) const {
return val==rhs.val;
}
};
vector<node> aim[1000];
ll acc[1000];

void modify_group(int idx,ll delta) {
int gid=idx/M;
vector<node> v1,v2;
ll pivot=-1;
bool added=false;
for (int i=0;i<aim[gid].size();i++) {
if (aim[gid][i].idx==idx) {
pivot=aim[gid][i].val-delta;
break;
}
}
for (int i=0;i<aim[gid].size();i++) {
if (aim[gid][i].idx<idx) {
if (!added&&pivot<=aim[gid][i].val) {
v1.pb({pivot,idx});
added=true;
}
v1.pb(aim[gid][i]);
} else if (aim[gid][i].idx>idx) {
v2.pb(aim[gid][i]);
v2.back().val+=delta;
}
}
if (!added) { v1.pb({pivot,idx}); }
aim[gid].clear();
merge(v1.begin(),v1.end(),v2.begin(),v2.end(),back_inserter(aim[gid]));
}

int main()
{
scanf("%d%d",&n,&q);
M=sqrt(n)+1;
for (int i=0;i<n;i++) {
scanf("%d",a+i);
s[i]=a[i];
if (i>0) s[i]+=s[i-1];
}
for (int i=0;i<n;i++) {
aim[i/M].pb({s[i]-2*a[i],i});
}
int g=n/M+1;
for (int i=0;i<g;i++) sort(aim[i].begin(),aim[i].end());
for (int i=0;i<q;i++) {
scanf("%d%d",&p,&x);
p--;
int id=p/M+1;
ll delta=x-a[p];
a[p]=x;
for (int j=id;j<g;j++) acc[j]+=delta;
modify_group(p,delta);
bool found=false;
for (int j=0;j<g;j++) {
node nd={-acc[j],-1};
auto iter=lower_bound(aim[j].begin(),aim[j].end(),nd);
if (iter!=aim[j].end()&&(*iter).val==-acc[j]) {
printf("%d\n",(*iter).idx+1);
found=true;
break;
}
}
if (!found) printf("-1\n");
}
}