Cod sursa(job #1975262)

Utilizator BogdanisarBurcea Bogdan Madalin Bogdanisar Data 30 aprilie 2017 13:03:23
Problema Arbori indexati binar Scor 100
Compilator cpp Status done
Runda Arhiva educationala Marime 3.3 kb
#include <iostream>
#include <fstream>

using namespace std;
ifstream in("aib.in");
ofstream out("aib.out");

typedef long long ll;
const int NMax = 1e5 + 5;

int N,M;
int aib[NMax];
// aib[i] = suma a[i - 2^r + 1] + a[i - 2^r + 2] + ... + a[i-1] + a[i],
// unde r este pozitia (numarata de la 0 si de la dreapta la stanga)
// ultimului bit (de la stanga la dreapta) cu valoarea 1 in reprezentarea binara a lui i

// nr & (~nr + 1) = o putere a lui 2 egala cu ultimul bit din reprezentarea binara a lui nr

void update(int,int);
int query(int);
int querySingle(int);
int findK(int);

int main() {
    in>>N>>M;

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

    while (M--) {
        int tip,a,b;
        in>>tip;
        if (!tip) {
            in>>a>>b;
            update(a,b);
        }
        else if (tip == 1) {
            in>>a>>b;

            if (a == b) {
                out<<querySingle(a)<<'\n';
            }
            else {
                out<<query(b) - query(a-1)<<'\n';
            }

        }
        else {
            in>>a;
            out<<( (a == 0) ? -1 : findK(a) )<<'\n';
        }
    }


    in.close();out.close();
    return 0;
}

// trebuie sa se faca update la toate pozitiile care au a[pozitia initiala] in suma lor;
// se poate arata ca daca se adauga ultimul bit la fiecare pos, se va obtine
// un index ce are a[pozitia initiala] in suma sa
// si ca acesta este cea mai mica valoare care se poate aduna la pos
// obtinandu-se astfel toate valorile de care avem nevoie
void update(int pos,int val) {
    while (pos <= N) {
        aib[pos] += val;

        pos += pos & (~pos + 1);
    }
}

// query(i) = suma a[1] + a[2] + ... + a[i]
// prin definitia lui aib[pos], se observa ca
// prefixSum[pos] = aib[pos] + prefixSum[ pos  - (pos & (~pos + 1)) ];
int query(int pos) {
    int sum = 0;
    while (pos) {
        sum += aib[pos];

        pos -= pos & (~pos + 1);
    }
    return sum;
}

// querySingle(i) = valoarea a[i];
// aib[pos] = a[lca + 1] + a[lca + 2] + ... + a[pos-1] + a[pos] (conform definitiei)
// sum2 = a[lca + 1] + a[lca + 2] + ... + a[pos-1] (  datorita reprezentarii binare a lui x-1
//                                                    si modului cum este construit aib-ul   )
// se gaseste raspunsul mai repede decat daca s-ar efectua query(pos) - query(pos -1);
int querySingle(int pos) {
    int sum1 = aib[pos],
        sum2 = 0,
        pre = pos - 1,
        lca = pos - (pos & (~pos + 1));
    while (pre != lca) {
        sum2 += aib[pre];

        pre -= pre & (~pre + 1);
    }
    return sum1 - sum2;
}

// findK(sum) = cel mai mare indice k pentru care a[1] + a[2] + ... + a[k] = sum;
// cel mai mic nu se poate gasi in logN, dar, fiindca toate numere a[i] sunt strict pozitive,
// suma a[1] + a[2] + ... + a[i], creste odata cu cresterea indicelui i,
// => exista un singur indice k care are suma sum (daca exista)
int findK(int sum) {
    int pw;
    for (pw = 1;pw <= N; pw <<= 1) ;
    pw >>= 1;

    int pos = 0;
    while (pw) {
        if (pos + pw <= N && aib[pos + pw] <= sum) {
            sum -= aib[pos += pw];
        }

        pw >>= 1;
    }

    if (!sum) {
        return pos;
    }

    return -1;
}