Cod sursa(job #2175486)

Utilizator Dddarius95Darius-Florentin Neatu Dddarius95 Data 16 martie 2018 17:27:43
Problema Arbori indexati binar Scor 100
Compilator cpp Status done
Runda Arhiva educationala Marime 1.71 kb
#include <bits/stdc++.h>
using namespace std;

#define LSB(x) ((-x) & x)
struct BIT {
	// UPDATE      UPDATE
	// PREFIX_SUM  SUM[1, ..., k]
	// QUERY 	   SUM[i, j]
	// FIND_K      SUM[1...k] == S

	
	int size;

	// because size is fixed, I can precompute STEP_MAX for FIND_K
	int STEP_MAX;
	vector<int> bit;

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

	// ~BIT() {
	// 	delete[] bit;
	// }

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

	int prefix_sum(int k) {
		int sum = 0;
		
		for(; k > 0; k -= LSB(k)) {
			sum += bit[k];
		}
		
		return sum;
	}
  	
  	int query(int i, int j) {
  		return prefix_sum(j) - prefix_sum(i - 1);
  	}
  	
  	int find_k(int sum) {
		int k = 0;
  		
  		for (int step(STEP_MAX); step > 0; step >>= 1) {
  			if (k + step <=  size && bit[k + step] <= sum) {
  				k += step;
  				sum -= bit[k];

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

  		return -1;
  	}
};

  
  
int main() {
	std::ios::sync_with_stdio(false); 
	ifstream fin("aib.in");
	ofstream fout("aib.out");

	int n, m;
	fin >> n >> m;

	BIT bit(n);

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

		// update
		bit.update(i, x);
	} 

	while (m--) {
		int type, start, end, x, pos;
		fin >> type;

		switch (type) {
			case 0: {
				fin >> pos >> x;
				bit.update(pos, x);

				break;
			}

			case 1: {
				fin >> start >> end;
				fout << bit.query(start, end) << "\n";

				break;
			}


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

				break;
			}

			default: {
				fout << "wtf?";
				break;
			}
		}
	}
	
	fin.close();
	fout.close();
	return 0;
}