AtCoder Library Practice Contest B - Fenwick Tree

BIT の使い方の練習問題。同時に、自分で BIT を書いたときにそれを verify できる問題でもある！

• 問題へのリンク

問題概要

長さ の数列 が与えられる。この数列に対する 個のクエリに答えよ。

• ：数列中の値 に を加算する (a[p] += x)
• ：数列 の区間 の総和を出力する

後者については、つまり、以下の式の値を出力する:

区間 [l, r) についての補足

数列 の区間 とは、 のことを指します。

区間 と書いたとき、左側は閉区間、右側は開区間となっています。つまり、 は区間に含まれますが、 は区間に含まれません。

制約

• 
• 

解法

BIT (Binary Indexed Tree / 別名 Fenwick Tree) は、まさに問題文で示された 2 つのクエリをサポートするデータ構造です。いずれも の計算量で実行できます。

前者の に を加算する処理については、一見 の計算量で実行できるようにも思えるのですが、後者の総和クエリに対応できるようにさまざまな値を書き換える必要があるため、 の計算量となります。

その代わりに、後者の一見 の計算量を要するようにも思える総和クエリも の計算量で答えられます。

BIT の仕組み自体を学びたい方には、たとえば次の記事に詳細が書かれています。

algo-logic.info

他にも、BIT で検索して有益な記事がたくさん出てきます。蟻本にも解説があります。

コード

BIT は、ACL では atcoder::fenwick_tree として実装されています。

#include <bits/stdc++.h>
#include <atcoder/fenwicktree>
using namespace std;

int main() {
    int N, Q;
    cin >> N >> Q;

    // BIT を宣言する (N 要素が 0 で初期化される)
    atcoder::fenwick_tree<long long> bit(N);

    // 数列 a を入力として受け取り、BIT に加算していく
    for (int i = 0; i < N; ++i) {
        long long a;
        cin >> a;
        bit.add(i, a);  // i 番目に a を加算
    }

    // 各クエリ
    for (int i = 0; i < Q; ++i) {
        int type;
        cin >> type;
        if (type == 0) {
            // 加算クエリ
            long long p, x;
            cin >> p >> x;
            bit.add(p, x);
        } else {
            // 総和取得クエリ
            int l, r;
            cin >> l >> r;

            // 区間 [l, r) の総和を取得して答える
            cout << bit.sum(l, r) << endl;
        }
    }
}

自前ライブラリでも AC

#include <bits/stdc++.h>
using namespace std;

template <class Abel> struct BIT {
    Abel UNITY_SUM = 0;
    vector<Abel> dat;
    
    // [0, n)
    BIT(int n, Abel unity = 0) : UNITY_SUM(unity), dat(n, unity) { }
    void init(int n) {
        dat.assign(n, UNITY_SUM);
    }
    
    // a is 0-indexed
    inline void add(int a, Abel x) {
        for (int i = a; i < (int)dat.size(); i |= i + 1)
            dat[i] = dat[i] + x;
    }
    
    // [0, a), a is 0-indexed
    inline Abel sum(int a) {
        Abel res = UNITY_SUM;
        for (int i = a - 1; i >= 0; i = (i & (i + 1)) - 1)
            res = res + dat[i];
        return res;
    }
    
    // [a, b), a and b are 0-indexed
    inline Abel sum(int a, int b) {
        return sum(b) - sum(a);
    }
    
    // debug
    void print() {
        for (int i = 0; i < (int)dat.size(); ++i)
            cout << sum(i, i + 1) << ",";
        cout << endl;
    }
};

int main() {
    int N, Q;
    cin >> N >> Q;

    // BIT を宣言する (N 要素が 0 で初期化される)
    BIT<long long> bit(N);

    // 数列 a を入力として受け取り、BIT に加算していく
    for (int i = 0; i < N; ++i) {
        long long a;
        cin >> a;
        bit.add(i, a);  // i 番目に a を加算
    }

    // 各クエリ
    for (int i = 0; i < Q; ++i) {
        int type;
        cin >> type;
        if (type == 0) {
            // 加算クエリ
            long long p, x;
            cin >> p >> x;
            bit.add(p, x);
        } else {
            // 総和取得クエリ
            int l, r;
            cin >> l >> r;

            // 区間 [l, r) の総和を取得して答える
            cout << bit.sum(l, r) << endl;
        }
    }
}