Cod sursa(job #3142362)

Utilizator NightCrawler92Alexandru Stefanica NightCrawler92 Data 20 iulie 2023 19:08:21
Problema Arbori indexati binar Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 1.63 kb
#include <vector>
#include <iostream>
#include <fstream>

class BinaryIndexTree {

public:
    using t = int;

    explicit BinaryIndexTree(size_t N): N_{N}, tree_(N + 1, 0) {} 

    void update(const t pos, const t value) {
        for (t i = pos; i <= N_;  i += lsb(i)) {
            tree_[i] += value;
        }
    }

    t query(const t pos) const {
        t sum = 0;
        for (auto i = pos; i >= 1; i -= lsb(i)) {
            sum += tree_[i];
        }
        return sum;
    }

    int findPosFor(const t sum) const {
        size_t idx = 0;
        size_t bitMask;

        t ssum = sum;
        for (bitMask = 1; bitMask <= N_; bitMask <<= 1);
        for (; bitMask; bitMask >>= 1) {
            if ( (idx + bitMask) <= N_ && tree_[idx + bitMask] < ssum) {
                idx += bitMask;
                ssum -= tree_[idx];
            }
        }

        if (query(idx + 1) != sum) {
            return -1;
        }
        return idx + 1;
    }


private:
    constexpr t lsb(t x) const {return (x & (-x));}

    const size_t N_;
    std::vector<t> tree_;

};

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

    int N, M;
    int op, a, b;

    in >> N >> M;
    BinaryIndexTree t(N);
    for (int i = 1; i <= N; ++i) {
        in >> a;
        t.update(i, a);
    }
    for (int i = 1; i <= M; ++i) {
        in >> op >> a;
        if (op == 2) {
            out << t.findPosFor(a) << '\n';
        } else {
            in >> b;
            if (op == 1) {
                out << ( t.query(b) - t.query(a - 1)) << '\n';
            } else {
                t.update(a, b);
            }
        }
    }

    return 0;


}