Cod sursa(job #798269)

Utilizator psycho21rAbabab psycho21r Data 16 octombrie 2012 00:26:29
Problema Arbori indexati binar Scor 100
Compilator cpp Status done
Runda Arhiva educationala Marime 1.86 kb
#include <fstream>
#include <vector>

using namespace std;

void update(int position, int quantity);
int query(int position);
int search(int a);
vector<int> sum;

int main()
{
    int N, M;
    ifstream in("aib.in");
    in >> N >> M;
    sum.resize(N + 1);

    for(int i = 1, temp_argument; i <= N; ++i)
    {
        in >> temp_argument;
        update(i, temp_argument);
    }

    ofstream out("aib.out");
    for(int i = 1, temp_type, a, b; i <= M; ++i)
    {
        in >> temp_type;
        switch(temp_type)
        {
            case 0:
                in >> a >> b;
                update(a, b);
                break;
            case 1:
                in >> a >> b;
                out << query(b) - query(a - 1) << "\n";
                break;
            case 2:
                in >> a;
                out << search(a) << "\n";
                break;
        }
    }
    out.close();

    return 0;
}

void update(int position, int quantity)
{
    int k = 0;
    while(position < sum.size())
    {
        sum[position] += quantity;
        while(!(position & (1 << k))) //Compute k
            ++k;
        position += (1 << k);
        k += 1;
    }
}

int query(int position)
{
    int k = 0, query_sum = 0;
    while(position > 0)
    {
        query_sum += sum[position];
        while(!(position & (1 << k))) //Compute k
            ++k;
        position -= (1 << k);
        k += 1;
    }
    return query_sum;
}

int search(int a)
{
    int step;
    for(step = 1; step < sum.size(); step <<= 1);
    for(int k = 0; step; step >>= 1)
    {
        if(k + step <= sum.size())
        {
            if(a >= sum[k + step])
            {
                k += step;
                a -= sum[k];
                if(!a)
                    return k;
            }
        }
    }
    return -1;
}