Cod sursa(job #2175416)

Utilizator rares96cheseliRares Cheseli rares96cheseli Data 16 martie 2018 17:04:27
Problema Arbori indexati binar Scor 30
Compilator cpp Status done
Runda Arhiva educationala Marime 2.37 kb
#include <bits/stdc++.h>
using namespace std;
 
#define LSB(x) (x & (-x))
struct BIT {
    // UPDATE      UPDATE
    // PREFIX_SUM  SUM[1, ..., k]
    // QUERY       SUM[i, j]
    // FIND_K      SUM[1...k] == S
 
     
    int size;
    vector<int> bit;
 
    // because size is fixed, I can precompute STEP_MAX for FIND_K
    int STEP_MAX;
 
    BIT(int _size) : size(_size), bit(size + 2, 0) {
         
        // precompute max_step
        STEP_MAX = 1;
        while (STEP_MAX < size) {
            STEP_MAX <<= 1;
        }
    }
 
    void update(int k, const int& x) {
        for(;  k <= size;  k += LSB(k)) {
            bit[k] += x;
        }
    }
 
    int prefix_sum(int k) {
        int sum = 0;
         
        for(; k > 0; k -= LSB(k)) {
            sum += bit[k];
        }
         
        return sum;
    }
     
    int query(int i, int j) {
        return prefix_sum(j) - prefix_sum(i - 1);
    }
     
    int find_k(int sum) {
        int k = 0;
         
        for (int step(STEP_MAX); step > 0; step >>= 1) {
            if (k + step <=  size && bit[k + step] <= sum) {
                k += step;
                sum -= bit[k];
 
                if (sum == 0) {
                    return k;
                }
            }
        }
 
        return -1;
    }
};
 
   
   
int main() {
    assert( freopen("aib.in", "r", stdin) != nullptr);
    assert( freopen("aib.out", "w", stdout) != nullptr);
    ios_base::sync_with_stdio(false);
 
    int n, m;
    cin >> n >> m;
 
    BIT bit(n);
 
    for (int i = 1; i <= n; ++i) {
        int x;
        cin >> x;
 
        // update
        bit.update(i, x);
    } 
 
    while (m--) {
        int type, start, end, x, pos;
        cin >> type;
 
        switch (type) {
            case 0: {
                cin >> pos >> x;
                bit.update(pos, x);
 
                break;
            }
 
            case 1: {
                cin >> start >> end;
                cout << bit.query(start, end) << "\n";
 
                break;
            }
 
 
            case 2: {
                cin >> x;
                cout << bit.find_k(x) << "\n";
 
                break;
            }
 
            default: {
                cout << "wtf?";
                break;
            }
        }
    }
     
    return 0;
}