Cod sursa(job #806251)

Utilizator dm1sevenDan Marius dm1seven Data 2 noiembrie 2012 14:28:47
Problema Arbori indexati binar Scor 20
Compilator cpp Status done
Runda Arhiva educationala Marime 4.3 kb
#include <iostream>
#include <fstream>
using namespace std;

//#include "../utils/PerformanceTimer.h"

struct Node
{
	int val;
	int left_id;
	int right_id;
	int left_sum;
	int right_sum;
};

void update_left_and_right_sum(int updated_id, Node* nodes, int id)
{
	if (updated_id < id)
	{
		//for left
		int left_id = nodes[id].left_id;
		update_left_and_right_sum(updated_id, nodes, left_id);
		nodes[id].left_sum = nodes[left_id].left_sum + nodes[left_id].val + nodes[left_id].right_sum;
	}
	
	if (updated_id > id)
	{
		//for right
		int right_id = nodes[id].right_id;
		update_left_and_right_sum(updated_id, nodes, right_id);
		nodes[id].right_sum = nodes[right_id].left_sum + nodes[right_id].val + nodes[right_id].right_sum;
	}
}

void compute_left_and_right_sum(Node* nodes, int id)
{
	if (id != 0)
	{
		int left_id = nodes[id].left_id;
		int right_id = nodes[id].right_id;
		
		//for left
		compute_left_and_right_sum(nodes, left_id);
		//for right
		compute_left_and_right_sum(nodes, right_id);

		//compute the current left_sum
		nodes[id].left_sum = nodes[left_id].left_sum + nodes[left_id].val + nodes[left_id].right_sum;
		//compute the current right_sum
		nodes[id].right_sum = nodes[right_id].left_sum + nodes[right_id].val + nodes[right_id].right_sum;
	}
}

int sum_first_k_nodes(int k, Node* nodes, int root_id)
{	
	int sum = 0;
	int current_id = root_id;
	while (current_id != k)
	{
		if (current_id < k)
		{
			sum += nodes[current_id].val + nodes[current_id].left_sum;
			current_id = nodes[current_id].right_id;//go to the right if k is higher than the current id
		}
		else current_id = nodes[current_id].left_id;		
	}
	//ends when current_id == k
	//update the sum with the current value and the current left sum
	sum += nodes[k].val + nodes[k].left_sum;
	
	return sum;
}

int index_for_sum(int sum, int s, Node* nodes, int id)//called with s = 0 and id = root_id
{
	//int s = 0;//the sum
	if (id == 0) return -1; //index not found

	int s1 = s + nodes[id].left_sum + nodes[id].val;
	
	if (s1 > sum) return index_for_sum(sum, s, nodes, nodes[id].left_id);
	if (s1 < sum) return index_for_sum(sum, s1, nodes, nodes[id].right_id);
	
	return id;//if s1 == sum
}

//int e_015_aib()
int main()
{
	//PerformanceTimer timer;
	//timer.init();

	char in_file[] = "aib.in";
	char out_file[] = "aib.out";

	const int MAX_Nodes = 2*100000;

	int N, M;

	ifstream ifs(in_file);
	ofstream ofs(out_file);
	ifs>>N>>M;
	
	int end_pow_2 = 1;
	while (end_pow_2 <= N) end_pow_2 *= 2;
	
	Node nodes[MAX_Nodes];	

	//build left_ids and right_ids
	nodes[0].val = 0;
	nodes[0].left_sum = 0;
	nodes[0].right_sum = 0;

	nodes[2].left_id = 1;
	nodes[2].right_id = 3;

	nodes[1].left_id = 0;
	nodes[1].right_id = 0;
	nodes[3].left_id = 0;
	nodes[3].right_id = 0;

	int pow_2 = 2;	
	while (2*pow_2 < end_pow_2)
	{	
		int old_pow_2 = pow_2;
		pow_2 *= 2;
		nodes[pow_2].left_id = old_pow_2;
		nodes[pow_2].right_id = pow_2 + old_pow_2;
		
		for (int i = 2; i < pow_2; i+=2)
		{
			nodes[i+pow_2].left_id = nodes[i].left_id + pow_2;//for even indices
			nodes[i+pow_2].right_id = nodes[i].right_id + pow_2;//for even indices			
		}
		for (int i = 1; i < pow_2; i+=2)
		{
			nodes[i + pow_2].left_id = 0;//for odd indices
			nodes[i + pow_2].right_id = 0;//for odd indices			
		}
	}


	for (int i = 1; i <= N; i++) ifs>>nodes[i].val;
	//fill with zeros the rest of nodes
	for (int i = N+1; i < end_pow_2; i++) nodes[i].val = 0;

	int root_id = end_pow_2/2;
	compute_left_and_right_sum(nodes, root_id);

	for (int m = 1; m <= M; m++) 
	{		
		int operation_id, a, b;
		ifs>>operation_id>>a;
		if (operation_id != 2)
		{
			ifs>>b;
		}

		int sum_a_1, sum_b;
		//execute the operation
		switch (operation_id)
		{
		case 0://operations_codes[0]:
			nodes[a].val += b;
			update_left_and_right_sum(a, nodes, root_id);
			break;
		case 1://operations_codes[1]:
			//for (int i = a; i <= b; i++) sum += A[i];
			sum_a_1 = sum_first_k_nodes(a-1, nodes, root_id);
			sum_b = sum_first_k_nodes(b, nodes, root_id);
			ofs<<sum_b - sum_a_1<<"\n";
			break;
		case 2://operations_codes[2]:
			ofs<<index_for_sum(a, 0, nodes, root_id)<<"\n";			
			break;
		}
	}	
	ifs.close();	
	ofs.close();	

	//cout<<timer.getElapsedTime()<<endl;

	return 0;
}