[AtCoder]ABC157 E – Simple String Queries (500点)

2020年3月2日AtCoder500点,セグメント木,区間,set,クエリ

問題概要

長さ N の英小文字からなる文字列 S を与えられ、以下の二種類のクエリ合計 Q 個を処理する。

  • S の \(i_q\) 番目を \(c_q\) に変更する
  • S の \(i_q\) 番目から \(r_q\) 番目までに出現する文字の種類を出力する

制約

  • \(1 \leq N \leq 500000\)
  • \(1 \leq Q \leq 20000\)

考え方

セグメント木を知っている場合

セグメント木というデータ構造を知っていれば

  • 変更のクエリと出力のクエリが入り乱れている
  • 区間に関係するクエリとなっている

という特徴に気がつくと、セグメント木が使えそうだと思うでしょう。実際に、一つのクエリに \(O(log N)\) 程度の計算量なら十分間に合います。

後はどうやって上手くセグメント木に落とし込むかという話に変わっていきます。

セグメント木を徹底解説！0から遅延評価やモノイドまで

セグメント木を使えるのはモノイドなので、どのようなモノイドならば問題の要求を満たせるかを考える必要があるでしょう。

今回は文字の種類を出力するので、「集合X：{a,b,c,…,z} の部分集合全て」「二項演算 \(\cup\)：\(a \cup b\)」などとすれば、

  • 結合則：\((a \cup b)\cup c =a \cup (b \cup c) \)
  • 単位元の存在：空集合 {}

が成立するので、これはモノイドです。

セグメント木を知らない場合

セグメント木を知らない場合でも解くことができます。

英小文字はたった 26 種類しか無いことから、「一つずつ分けて考える」ことができそうです。

一つのクエリあたり、\(O(log N)\) 程度で処理しなければならないことを考えると、文字ごとに考えることにすると更新のクエリは消去される文字と新しい文字の2つしか処理を行いません。

  • 更新処理（ある場所から文字を消去 or ある場所に文字を新規に挿入）に \(2 \times O(log N)\) 程度

区間取得のクエリは、それぞれの文字ごとに、区間内に1つでも出現するものがあるかどうかを判定すればよいです。

  • 区間取得（ある区間にその文字が出現するかどうか判定）に \(26 \times O(log N)\) 程度

以上を考えると、英小文字の種類ごとに、出現位置を記録した set(平衡二分探索木) を使えば上手くいくことが分かります。

解法（セグメント木）

以下のモノイドをセグメント木に乗せて、更新と区間取得のクエリに答えればよいです。

  • 集合X：{a,b,c,…,z} の部分集合全て
  • 二項演算 \(\cup\)：\(a \cup b\)

実装上は、26種類の文字のどれが出現したかどうかを管理する方法は色々あり、

  • set(平衡二分探索木) で管理
  • 配列(or vector) で管理
  • ビットで管理

などが考えられますが、ビットで管理すれば二項演算 \(\cup\) は or 演算子を使用することができるので高速に動作します。

C++での実装例

ビットの立っている数は __builtin_popcount 関数などで取得することができます。セグメント木についてはこちらを参考のこと。

#include <bits/stdc++.h>
#define LOOP(n) for (int _i = 0; _i < (n); _i++)
#define REP(i, n) for (int i = 0; i < (n); ++i)
#define RREP(i, n) for (int i = (n); i >= 0; --i)
#define FOR(i, r, n) for (int i = (r); i < (n); ++i)
#define ALL(obj) begin(obj), end(obj)
using namespace std;
/* SegTree<X>(n,fx,ex): モノイド(集合X, 二項演算fx, 単位元ex)についてサイズnで構築
    set(int i, X x), build(): i番目の要素をxにセット。まとめてセグ木を構築する。O(n)
    update(i,x): i 番目の要素を x に更新。O(log(n))
    query(a,b):  [a,b) 全てにfxを作用させた値を取得。O(log(n))
*/
template <typename X>
struct SegTree {
    using FX = function<X(X, X)>;
    int n;
    FX fx;
    const X ex;
    vector<X> dat;
    SegTree(int n_, FX fx_, X ex_) : n(), fx(fx_), ex(ex_), dat(n_ * 4, ex_) {
        int x = 1;
        while (n_ > x) {
            x *= 2;
        }
        n = x;
    }

    void set(int i, X x) { dat[i + n - 1] = x; }
    void build() {
        for (int k = n - 2; k >= 0; k--) dat[k] = fx(dat[2 * k + 1], dat[2 * k + 2]);
    }

    void update(int i, X x) {
        i += n - 1;
        dat[i] = x;
        while (i > 0) {
            i = (i - 1) / 2;  // parent
            dat[i] = fx(dat[i * 2 + 1], dat[i * 2 + 2]);
        }
    }

    // the minimum element of [a,b)
    X query(int a, int b) { return query_sub(a, b, 0, 0, n); }
    X query_sub(int a, int b, int k, int l, int r) {
        if (r <= a || b <= l) {
            return ex;
        } else if (a <= l && r <= b) {
            return dat[k];
        } else {
            X vl = query_sub(a, b, k * 2 + 1, l, (l + r) / 2);
            X vr = query_sub(a, b, k * 2 + 2, (l + r) / 2, r);
            return fx(vl, vr);
        }
    }
    /* debug */
    inline X operator[](int a) { return query(a, a + 1); }
    void print() {
        for (int i = 0; i < n; ++i) {
            cout << (*this)[i];
            if (i != n) cout << ",";
        }
        cout << endl;
    }
};

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    int N, Q;
    string S;
    cin >> N >> S >> Q;
    using X = int;
    auto fx = [](X x1, X x2) -> X { return x1 | x2; };
    X ex = 0;
    SegTree<X> seg(N, fx, ex);

    REP(i, N) {
        X x = 1 << (int)(S[i] - 'a');
        seg.set(i, x);
    }
    seg.build();

    vector<int> ans;

    for (int i = 0; i < Q; i++) {
        int q;
        cin >> q;
        if (q == 1) {
            int i;
            char c;
            cin >> i >> c;
            i--;
            X x = 1 << (int)(c - 'a');
            seg.update(i, x);
        } else if (q == 2) {
            int l, r;
            cin >> l >> r;
            X ans = seg.query(l - 1, r);
            cout << __builtin_popcount(ans) << "\n";
        }
    }

    return 0;
}

解法（set を使う場合）

26種の文字ごとに、出現位置を記録した set を用意します。すると、以下のようにクエリに対して set のメソッドが利用できます。

  • 更新処理（ある場所から文字を消去 or ある場所に文字を新規に挿入）に set の erase や insert
  • 区間取得（ある区間にその文字が出現するかどうか判定）に set の lower_bound（l 以降の出現位置が r 以下かどうか判定）

どちらも1回あたり \(O(log N)\) 程度で計算できるので十分高速です。

C++ での実装例

#include <bits/stdc++.h>
#define LOOP(n) for (int _i = 0; _i < (n); _i++)
#define REP(i, n) for (int i = 0; i < (n); ++i)
#define RREP(i, n) for (int i = (n); i >= 0; --i)
#define FOR(i, r, n) for (int i = (r); i < (n); ++i)
#define ALL(obj) begin(obj), end(obj)
using namespace std;

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

    set<int> s[26];  // 0:a, 1:b, 2:c, ..., 25:z
    REP(i, N) { s[(int)(S[i] - 'a')].insert(i); }

    vector<int> ans;
    for (int i = 0; i < Q; i++) {
        int q;
        cin >> q;
        if (q == 1) {
            int i;
            char c;
            cin >> i >> c;
            i--;
            s[(int)(S[i] - 'a')].erase(i);
            s[(int)(c - 'a')].insert(i);
            S[i] = c;
        } else if (q == 2) {
            int l, r;
            cin >> l >> r;
            l--, r--;
            int sum = 0;
            REP(i, 26) {
                auto it = s[i].lower_bound(l);
                if (it != s[i].end()) {
                    if (*it <= r) {
                        sum++;
                    }
                }
            }

            ans.push_back(sum);
        }
    }

    for (auto i : ans) {
        cout << i << endl;
    }

    return 0;
}

AtCoder500点,セグメント木,区間,set,クエリ