Cod sursa(job #1759579)

Utilizator andreioneaAndrei Onea andreionea Data 19 septembrie 2016 16:05:10
Problema Arbori indexati binar Scor 100
Compilator cpp Status done
Runda Arhiva educationala Marime 1.32 kb
#include <fstream>
#include <vector>

class AIB
{
public:
	AIB(int);
	void Insert(int, std::size_t);
	int GetSum(int, int) const;
	int GetSum(int) const;
	int GetPos(int) const;
private:
	std::vector<int> v;
};

AIB::AIB(int n): v(n, 0) {}


void AIB::Insert(int val, std::size_t pos)
{
	if (pos > v.size())
		return;
	v[pos - 1] += val;
	Insert(val, pos + (pos & (-pos)));
}

int AIB::GetSum(int pos) const
{
	if (pos <= 0)
		return 0;
	if (pos == 1)
		return v[pos - 1];
	return v[pos - 1] + GetSum(pos - (pos & (-pos)));
}

int AIB::GetSum(int posA, int posB) const
{
	int a = GetSum(posA - 1), b = GetSum(posB);
	return b - a;
}

int AIB::GetPos(int a) const
{
	int fst = 1;
	int lst = v.size();
	while (fst <= lst) {
		auto mid = (fst + lst) / 2;
		auto sum = GetSum(mid);
		if (sum < a)
			fst = mid + 1;
		else if (sum > a)
			lst = mid - 1;
		else {
			return mid;
		}
	}
	return -1;
}

int main()
{
	int n, m;
	std::ifstream f("aib.in");
	std::ofstream g("aib.out");
	int op;
	int a,b;
	f >> n >> m;
	AIB aib(n);
	for(int i = 1; i <= n; ++i){
		int x;
		f >> x;
		aib.Insert(x, i);
	}
	while (m--) {
		f >> op;
		if (op == 0) {
			f >> a >> b;
			aib.Insert(b, a);
		}
		else 
		if (op == 1) {
			f >> a >> b;
			g << aib.GetSum(a, b) << "\n";
		}
		else {
			f >> a;
			g << aib.GetPos(a) << "\n";
		}
	}
	f.close();
	g.close();
	return 0;
}