IOICJ 36 Interval Sum EX ( Math, Segment Tree )

Problem Description:
You are given an array A wit length N.
You need to support these operations, noted with [ p, e, c ]:
if p = 0:
__square root and discard non-integer value of range [ e, c ]
if p = 1:
__change value at e to c
if p = 2:
__output the sum of [ e, c ]

Constraints:
1≤T≤10
1≤N,M≤1e5
1≤Ai≤1e9
Sqrting and Querying: 1≤e≤c≤N
Modification: 1≤e≤N,1≤c≤1e9

Sample Input:
1
4 10
4 8 16 27
2 1 3
0 2 4
2 3 4
1 2 120
2 1 4
0 1 3
2 1 2
1 1 1
0 1 4
2 2 4

Sample Output:
28
9
133
12
6

Solution:
Since all values are below 1e9, and modification could only be applied to a single point each time, we know even if all values were initially 1e9 and all modifications are set to 1e9, and we are to square root on every single element, at most ( N + Q ) * lg( 1e9 ) operations will be performed ( after that all values will be 1 ). Therefore on sum-maintaining segment tree, we would additionally record whether all values in that range is 1.

#include <bits/stdc++.h>
using namespace std;

#define int long long

const int MAXN = ( int ) 1e5;

int N, Q;
int A[ MAXN ];

int segt_sum[ MAXN << 3 ];

void segt_pull( int t ){
  segt_sum[ t ] = segt_sum[ t << 1 ] + segt_sum[ t << 1 | 1 ];
}

void segt_init( int t, int *a, int lb, int rb ){ // [ lb, rb )
  if( lb + 1 == rb ){
    segt_sum[ t ] = a[ lb ];
    return;
  }
  int mid = lb + rb >> 1;
  segt_init( t << 1, a, lb, mid );
  segt_init( t << 1 | 1, a, mid, rb );
  segt_sum[ t ] = segt_sum[ t << 1 ] + segt_sum[ t << 1 | 1 ];
}

void update( int t, int lb, int rb, int ql, int qr ){
  if( qr <= lb or rb <= ql ) return;
  if( lb + 1 == rb ){
    segt_sum[ t ] = sqrt( segt_sum[ t ] );
    return;
  }
  if( segt_sum[ t ] == rb - lb ) return;
  int mid = lb + rb >> 1;
  update( t << 1, lb, mid, ql, qr );
  update( t << 1 | 1, mid, rb, ql, qr );
  segt_pull( t );
}

void modify( int t, int lb, int rb, int k, int v ){
  if( lb + 1 == rb ){
    segt_sum[ t ] = v;
    return;
  }
  int mid = lb + rb >> 1;
  if( k < mid ) modify( t << 1, lb, mid, k, v );
  else modify( t << 1 | 1, mid, rb, k, v );
  segt_pull( t );
}

int query( int t, int lb, int rb, int ql, int qr ){
  if( qr <= lb or rb <= ql ) return 0;
  if( ql <= lb and rb <= qr ){
    return segt_sum[ t ];
  }
  int mid = lb + rb >> 1;
  return query( t << 1, lb, mid, ql, qr ) + query( t << 1 | 1, mid, rb, ql, qr );
}

signed main(){
  ios::sync_with_stdio( 0 );
  int T; cin >> T;
  for( int ti = 0; ti < T; ++ti ){
    cin >> N >> Q;
    for( int i = 0; i < N; ++i ){
      cin >> A[ i ];
    }
    segt_init( 1, A, 0, N );
    for( int i = 0; i < Q; ++i ){
      int op; cin >> op;
      if( op == 0 ){ // do sqrt to [ ql, qr ]
        int ql, qr; cin >> ql >> qr;
        update( 1, 0, N, ql - 1, qr );
      } else if( op == 1 ){ // modify [ x, x ] to c
        int x, c; cin >> x >> c;
        modify( 1, 0, N, x - 1, c );
      } else{ // print sum of [ ql, qr ]
        int ql, qr; cin >> ql >> qr;
        cout << query( 1, 0, N, ql - 1, qr ) << endl;
      }
    }
  }
  return 0;
}