Дерево Фенвика с модификацией на отрезке

С этой структурой данных можно ознакомиться в этом посте и её модификацией для нахождения максимума в этом. Но я нигде не встречал реализацию с изменением элементов на отрезке, поэтому решил поделиться тем, что сумел получить самостоятельно.

1. Постановка задачи

Имеется массив. Необходимо уметь выполнять запросы двух видов:

1. Модификация — прибавить d к каждому элементу отрезка [l, r]
2. Сумма — вычислить сумму элементов отрезка [l, r]

2. Описание решения

Нетрудно заметить, что оба вида запросов можно «облегчить» до отрезка [0, r], разбивая отрезок [l, r] на два префикса. Как и для дерева отрезков, заведём второй массив, в котором будем хранить сколько надо прибавить ко всем числам этого отрезка.

Создадим класс Fenwick с прототипами будущих методов.

class Fenwick{
  int *m, *mt, N;
public:
  Fenwick(int n);
  Fenwick(int a[], int n);
  void add_range(int r, int d);
  void add_range(int l, int r, int d);
  int sum(int r);
  int sum(int l, int r);
};

В первом конструкторе достаточно выделить память и обнулить элементы массива. А во втором ещё пройдёмся по массиву a и прорелаксируем значения в дереве.

Fenwick::Fenwick(int n){
  N = n;
  m = new int[N];
  mt = new int[N];
  memset(m, 0, sizeof(int)*N);
  memset(mt, 0, sizeof(int)*N);
}
Fenwick::Fenwick(int a[], int n){
  N = n;
  m = new int[N];
  mt = new int[N];
  memset(m, 0, sizeof(int)*N);
  memset(mt, 0, sizeof(int)*N);
  for(int i=0;i<N;++i){
    m[i] += a[i];
    if((i|(i+1))<N) m[i|(i+1)] += m[i];
  }
}

Рассмотрим теперь операцию изменения. Предположим пришёл запрос «прибавить 1 к элементам отрезка [0, 9]». Этот отрезок полностью покрывается непересекающимися отрезками [0, 7] и [8, 9], поэтому в 7 и 9 элементы массива mt прибавляем 1. Но есть отрезки, содержащие [0, 9] (или его часть) в качестве подотрезка. Все они располагаются справа. Для них необходимо изменить значение в массиве m на столько, сколько элементов отрезка [0, 9] они содержат. То есть в m[11] прибавить 2, а к m[15] — 10.

Из рисунка видно, что перемещения происходят по тем же законам, что и в тривиальной реализации дерева Фенвика, поэтому сразу перейдём к коду.

void Fenwick::add_range(int r, int d){
  if(r<0) return ;
  for(int i=r;i>=0;i=(i&(i+1))-1) mt[i] += d;
  for(int i=r|(r+1);i<N;i|=i+1) m[i] += d*(r-(i&(i+1))+1);
}
void Fenwick::add_range(int l, int r, int d){
  add_range(r, d);
  add_range(l-1, -d);
}

Для суммы на префиксе подойдёт та же картинка с тем лишь пояснением, что для зелёных элементов необходимо прибавить и m, и mt, а для синих только mt (то есть учесть большие накрывающие отрезки).

int Fenwick::sum(int r){
  if(r<0) return 0;
  int res = 0;
  for(int i=r;i>=0;i=(i&(i+1))-1) res += m[i] + mt[i]*(i-(i&(i+1))+1);
  for(int i=r|(r+1);i<N;i|=i+1) res += mt[i]*(r-(i&(i+1))+1);
  return res;
}
int Fenwick::sum(int l, int r){
  return sum(r) - sum(l-1);
}

3. Итоги

Нам удалось получить структуру данных с инициализацией за O(N), модификацией и суммой на отрезке за O(logN). На рандомном тесте с 10000000 элементов и 10000000 запросов получил такие цифры:

Автор: kormyshov

Источник