Cod sursa(job #2704282)

Utilizator vnedelcuVictor Andrei Nedelcu vnedelcu Data 10 februarie 2021 10:15:35
Problema Arbori indexati binar Scor 50
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.16 kb
#include <stdio.h>
#include <vector>

using namespace std;

class FenwickTree {
private:
    int no_nodes;
    vector<int> nodes;
    int get_segment_size(int node) {
        return ((node ^ (node - 1)) & node);
    }

public:
    FenwickTree(int no_nodes);
    void update_node(int node, int value);
    int get_sum_to_node(int node);
    int find_sum_occurence(int sum);
};

FenwickTree::FenwickTree(int no_nodes) {
    this->no_nodes = no_nodes;
    nodes.resize(no_nodes + 1);
}

void FenwickTree::update_node(int node, int value) {
    while (node <= no_nodes) {
        nodes[node] += value;
        node += get_segment_size(node);
    }
}

int FenwickTree::get_sum_to_node(int node) {
    int sum = 0;
    while (node > 0) {
        sum += nodes[node];
        node -= get_segment_size(node);
    }
    return sum;
}

int FenwickTree::find_sum_occurence(int sum) {
    int index = 0;

    int bit = 1;
    for (; bit < no_nodes; bit <<= 1);

    for (; bit > 0; bit >>= 1) {
        if (index + bit <= no_nodes && nodes[index + bit] <= sum) {
            index += bit;
            sum -= nodes[index];
        }
    }

    if (sum != 0) {
        index = -1;
    }

    return index;
}

int main() {
    FILE * fin, * fout;
    int n, m, op, val, a, b;

    fin = fopen("aib.in", "r");

    fscanf(fin, "%d%d", &n, &m);
    FenwickTree fenwick_tree(n);

    for (int i = 1; i <= n; i++) {
        fscanf(fin, "%d", &val);
        fenwick_tree.update_node(i, val);
    }

    fout = fopen("aib.out", "w");
    for (int i = 0; i < m; i++) {
        fscanf(fin, "%d", &op);

        if (op == 0) {
            fscanf(fin, "%d%d", &a, &b);
            fenwick_tree.update_node(a, b);
        } else if (op == 1) {
            fscanf(fin, "%d%d", &a, &b);
            val = fenwick_tree.get_sum_to_node(b) - 
                fenwick_tree.get_sum_to_node(a - 1);
            fprintf(fout, "%d\n", val);
        } else {
            fscanf(fin, "%d", &a);
            val = fenwick_tree.find_sum_occurence(a);
            fprintf(fout, "%d\n", val);
        }
    }

    fclose(fin);
    fclose(fout);

    return 0;
}