인덱스 트리

C++

#include <cstdio>
#include <cstring>

#define MINIMUM_POWER_OF_2_GREATER_THAN_MAXIMUM_N 131072

int T[MINIMUM_POWER_OF_2_GREATER_THAN_MAXIMUM_N*2];
int M = 1;

void init(int N) {
    while (M < N) {
        M *= 2;
    }
}

void update(int index, int delta) {
    index = M + index - 1;
    while (index > 0) {
        T[index] += delta;
        index /= 2;
    }
}

int sum(int left, int right) {
    left = left + M - 1; right = right + M - 1;
    int sum = 0;
    while (left <= right) {
        if ((left & 1) == 1) sum += T[left++];
        if ((right & 1) == 0) sum += T[right--];
        left /= 2;
        right /= 2;
    }
    return sum;
}

Swift

import Foundation

struct IndexedTree<Element> {
  private var storage: Array<Element?>
  private var capacity: Int
  private var operation: (Element?, Element?) -> Element?
  private var firstLeafIndex: Int
  private var size: Int = 0

  init(capacity: Int, elements: [Element], operation: @escaping (Element, Element) -> Element) {
    self.capacity = capacity
    self.operation = { lhs, rhs in
      if let lhs = lhs, let rhs = rhs {
        return operation(lhs, rhs)
      } else if let lhs = lhs {
        return lhs
      } else if let rhs = rhs {
        return rhs
      } else {
        return nil
      }
    }
    firstLeafIndex = 1
    while firstLeafIndex < capacity {
      firstLeafIndex *= 2
    }
    storage = Array<Element?>(repeating: nil, count: firstLeafIndex * 2)
    elements.map { add($0) }
  }
  
  private mutating func add(_ element: Element) {
    var index = size + firstLeafIndex
    while index > 0 {
      print(index)
      storage[index] = operation(storage[index], element)
      index /= 2
    }
    size += 1
  }
  
  func idealValue(_ lowerBound: Int, _ upperBound: Int) -> Element? {
    guard 0 <= lowerBound, upperBound < size else { return nil }
    var lower = lowerBound + firstLeafIndex
    var upper = upperBound + firstLeafIndex
    var idealValue: Element?
    while lower < upper {
      if (lower & 1) == 1, let lowerValue = storage[lower] {
        idealValue = operation(idealValue, lowerValue)
        lower += 1
      }
      if (upper & 1) == 0, let upperValue = storage[upper] {
        idealValue = operation(idealValue, upperValue)
        upper -= 1
      }
      lower /= 2
      upper /= 2
    }
    if lower == upper, let lowerValue = storage[lower] {
      idealValue = operation(idealValue, lowerValue)
    }
    return idealValue
  }
}