Cod sursa(job #2221949)

Utilizator claudiu.gatinaFMI Claudiu Gatina claudiu.gatina Data 16 iulie 2018 10:41:27
Problema Arbori indexati binar Scor 0
Compilator cpp Status done
Runda Arhiva educationala Marime 2.39 kb
#include <iostream>
#include <cstdio>
#include <algorithm>

using namespace std;

class aib
{
  int arb[200010];
  int n;
  int h;
  void transformToTreeIdx(int &position)
  {
    position = (1 << h) + position - 1;
  }
  void calcHeight()
  {
    int r = n - 1;
    h = 0;
    while(r)
    {
      ++h;
      r >>= 1;
    }
  }
public:
  void add(int position, int val)
  {
    transformToTreeIdx(position);
    while(position)
    {
      arb[position] += val;
      position >>= 1;
    }
  }
  void init(int _n)
  {
    n = _n;
    for(int i = 0; i <= n << 1; ++i)
    {
      arb[i] = 0;
    }
    calcHeight();
    for(int i = 1; i <= n; ++i)
    {
      int x;
      scanf("%d", &x);
      add(i, x);
    }
  }
  int intervalSum(int l, int r)
  {
    int sub = 0;
    transformToTreeIdx(l);
    transformToTreeIdx(r);
    while(l != r)
    {
      if(l & 1)
        sub += arb[l - 1];
      if(!(r & 1))
        sub += arb[r + 1];
      l >>= 1;
      r >>= 1;
    }
    return arb[l] - sub;
  }
  int findk(int s)
  {
    int k = 1;
    transformToTreeIdx(k);

    while(arb[k] < s && k)
    {
      k >>= 1;
    }
    if(!k)
      return -1;
    else
    {
      int sum = arb[k];
      while(k < 1 << h)
      {
        if(sum - arb[(k << 1) + 1] <= s)
        {
          k <<= 1;
          ++k;
        }
        else
        {
          k <<= 1;
          sum -= arb[k + 1];
        }

      }
      if(sum == s)
        return min(n, k - (1 << h) + 1);
      else
        return -1;
    }
  }
};

int main()
{
    freopen("aib.in", "r", stdin);
    freopen("aib.out", "w", stdout);
    aib arb;
    int n, m;
    scanf("%d", &n);
    scanf("%d", &m);
    arb.init(n);
    for(int i = 0; i < m; ++i)
    {
      int op;
      scanf("%d", &op);
      switch(op)
      {
      case 0:
        {
          int a, b;
          scanf("%d", &a);
          scanf("%d", &b);
          arb.add(a, b);
        }
        break;
      case 1:
        {
          int a, b;
          scanf("%d", &a);
          scanf("%d", &b);
          printf("%d\n", arb.intervalSum(a, b));
        }
        
        break;
      case 2:
        {
          int a;
          scanf("%d", &a);
          printf("%d\n", arb.findk(a));
        }
        break;
      default:
        break;
      }
    }
    return 0;
}