Cod sursa(job #941849)

Utilizator sebii_cSebastian Claici sebii_c Data 19 aprilie 2013 21:41:38
Problema Arbori indexati binar Scor 60
Compilator cpp Status done
Runda Arhiva educationala Marime 1.9 kb
#include <fstream>
#include <vector>

class fenwick_tree {
public:
    explicit fenwick_tree(const std::vector<int>& vals);

    int query(int id);
    int query(int left, int right);
    int find(int sum);
    void update(int id, int to_add);

private:
    std::vector<int> tree;
};

fenwick_tree::fenwick_tree(const std::vector<int>& vals) {
    tree.resize(vals.size() + 1);
    for (int i = 0; i < vals.size(); ++i) 
        update(i + 1, vals[i]); 
}

int fenwick_tree::query(int id) {
    int sum = 0;
    while (id > 0) {
        sum += tree[id];
        id -= (id & -id);
    }

    return sum;
}

int fenwick_tree::query(int left, int right) {
    return query(right) - query(left - 1);
}

void fenwick_tree::update(int id, int to_add) {
    while (id < tree.size()) {
        tree[id] += to_add;
        id += (id & -id);
    }
}

int fenwick_tree::find(int sum) {
    int bitmask;
    for (bitmask = 1; bitmask < tree.size() - 1; bitmask *= 2);

    for (int i = 0; bitmask != 0; bitmask /= 2) {
        if (i + bitmask < tree.size()) {
            if (sum >= tree[i + bitmask]) {
                i += bitmask;
                sum -= tree[i];
                if (sum == 0)
                    return i;
            }
        }
    }

    return -1;
}

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

    int N, M;
    fin >> N >> M;

    std::vector<int> vals(N);
    for (int i = 0; i < N; ++i)
        fin >> vals[i];

    fenwick_tree tree(vals);
    for (int i = 0; i < M; ++i) {
        int op;
        fin >> op;

        if (op == 0) {
            int id, val;
            fin >> id >> val;
            tree.update(id, val);
        } else if (op == 1) {
            int l, r;
            fin >> l >> r;
            fout << tree.query(l, r) << "\n";
        } else {
            int sum;
            fin >> sum;
            fout << tree.find(sum) << "\n";
        }
    }

    return 0;
}