Range Sum Query and Range Add Query (RSQ and RAQ)

概要

以下のクエリが処理できるBinary Indexed Tree実装

  • \(a_l, a_{l+1}, ..., a_{r-1}\) に \(x\) を加算

  • \(a_l, a_{l+1}, ..., a_{r-1}\) の和を求める

具体的計算

蟻本ベースの説明[1]で、Binary Indexed Treeを2つ用いた実装。

\(P_i = p_1 + ... + p_i\), \(Q_i = q_1 + ... + q_i\) として、 \(a_1, ..., a_i\) の和 \(S_i\) を \(P_i * i + Q_i\) と表すことを考える。

こうすると、\(a_l, ..., a_{r-1}\)の和は \(S_{r-1} - S_{l-1}\) で計算できる。

この時、\(a_l, ..., a_{r-1}\) に \(x\) を加算したい時、以下の演算により更新できる。

  • \(q_l \leftarrow q_l - x(l-1)\)

  • \(q_r \leftarrow q_r + x(r-1)\)

  • \(p_l \leftarrow p_l + x\)

  • \(p_r \leftarrow p_r - x\)

そして、\(P_i, Q_i\)の値をそれぞれBITで管理することで、加算と和計算の1クエリの計算量を \(O(\log N)\) にすることができる。

計算量

  • 区間加算: \(O(\log N)\)

  • 区間総和: \(O(\log N)\)

実装

# N: 処理する区間の長さ

data0 = [0]*(N+1)
data1 = [0]*(N+1)
# 区間[l, r)に x を加算
def _add(data, k, x):
    while k <= N:
        data[k] += x
        k += k & -k
def add(l, r, x):
    _add(data0, l, -x*(l-1))
    _add(data0, r, x*(r-1))
    _add(data1, l, x)
    _add(data1, r, -x)

# 区間[l, r)の和を求める
def _get(data, k):
    s = 0
    while k:
        s += data[k]
        k -= k & -k
    return s
def query(l, r):
    return _get(data1, r-1) * (r-1) + _get(data0, r-1) - _get(data1, l-1) * (l-1) - _get(data0, l-1)

Verified

  • AOJ: "DSL_2_G: Range Query - RSQ and RAQ": source (Python3, 1.33sec)



1. プログラミングコンテストチャレンジブック [第2版] p.163