Cod sursa(job #2203598)

Utilizator Menage_a_011UPB Cheseli Neatu Popescu Menage_a_011 Data 12 mai 2018 18:36:36
Problema Arbori indexati binar Scor 100
Compilator cpp Status done
Runda Arhiva educationala Marime 2.36 kb
// https://goo.gl/fBmFxu
#include <bits/stdc++.h>
using namespace std;

#define NMAX        100009
#define MMAX        200009
#define kInf        (1 << 30)
#define kInfLL      (1LL << 60)
#define kMod        666013
#define edge pair<int, int>
#define x first
#define y second

#define USE_FILES "MLC"

#ifdef USE_FILES
#define cin fin
#define cout fout
ifstream fin("aib.in");
ofstream fout("aib.out");
#endif

// number of tests from "in"
int test_cnt = 1;
void clean_test();

// your global variables are here

// Binary Indexed Trees
template<typename T = int>
struct AIB {
	// UPDATE      UPDATE
	// PREFIX_SUM  SUM[1, ..., k]
	// QUERY 	   SUM[i, j]
	// FIND_K      SUM[1...k] == S
	#define LSB(x) ((-x) & x)

	int size;
	vector<T> aib; // aib[0], aib[1], ..., aib[size]

	// because size is fixed, I can precompute STEP_MAX for FIND_K
	int STEP_MAX;

	AIB(int _size) : size(_size), aib(_size + 1, 0) {
		// precompute max_step
		STEP_MAX = 1;
  		while (STEP_MAX < size) {
  			STEP_MAX <<= 1;
  		}
	}

	void update(int k, const T& x) {
		for(;  k <= size;  k += LSB(k)) {
			aib[k] += x;
		}
	}

	int prefix_sum(int k) {
		T sum = 0;

		for(; k > 0; k -= LSB(k)) {
			sum += aib[k];
		}

		return sum;
	}

  	T query(const int &i, const int &j) {
  		return prefix_sum(j) - prefix_sum(i - 1);
  	}

  	int find_k(T sum) {
		int k = 0;

  		for (int step(STEP_MAX); step > 0; step >>= 1) {
  			if (k + step <=  size && aib[k + step] <= sum) {
  				k += step;
  				sum -= aib[k];

  				if (sum == 0) {
  					return k;
  				}
  			}
  		}

  		return -1;
  	}
};

// your solution is here
void solve() {
    int n, Q;
    cin >> n  >> Q;

    AIB<int> aib(n + 1);
    vector<int> v(n + 1);

    for (int i = 1; i <= n; ++i) {
        int x;
        cin >> x;

        aib.update(i, x);
    }

	while (Q--) {
		int type, start, end, x, pos;
		cin >> type;

		switch (type) {
			case 0: {
				cin >> pos >> x;
				aib.update(pos, x);
				v[pos] = x;

				break;
			}

			case 1: {
				cin >> start >> end;
				cout << aib.query(start, end) << "\n";

				break;
			}


			case 2: {
				fin >> x;
				fout << aib.find_k(x) << "\n";

				break;
			}

			default: {
				fout << "wtf?";
				break;
			}
		}
	}


    if (test_cnt > 0) {
        clean_test();
    }
}


void clean_test() {
    // clean if needed
}


int main() {
//    cin >> test_cnt;
    while (test_cnt--) {
        solve();
    }

    return 0;
}