けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 157 E - Simple String Queries (水色, 500 点)

一瞬、「更新がある」「種類数」ということで、Wavelet Matrix が必要かとびびった

問題へのリンク

問題概要

長さ  N の文字列  S が与えられる。以下の  Q 個のクエリに答えよ

  • タイプ 1:  S i 文字目を  c に変更する
  • タイプ 2:  S の区間  \lbrack l, r \rbrack に含まれる文字の種類数を答える

制約

  •  N, Q \le 10^{5}

考えたこと

もしタイプ 2 だけだったら、種類数は、以下の記事みたいな方法で求めることができる。

drken1215.hatenablog.com

しかし更新付きだとそうはいかない。一般には

  • 二次元セグメントツリー
  • 動的 Wavelet Matrix

とか、そういうのが必要になってしまう。しかしこの問題は 500 点問題だ。そんなものが必要なはずがない。この問題ならではの特徴として、以下のものがある

  • 登場する文字の種類は多くても 26 種類しかない!!!

それを活かすと単純に、区間 [left, right) が与えられたときに、各文字 c について区間内に含まれるかどうかを判定する、という方針がとれる。この判定が十分高速にできるなら、それを 26 倍した計算時間で 1 クエリに対応できるのだ。そのための方法として、こんなものがある

  1. set を使う (とても簡単)
  2. BIT を 26 個使う

解法 (1): set を使う

各文字 c につき「文字が c であるような index の集合を表す set」を用意する (その変数名を a とかにする)。そうすると

  • i 文字目を p から c に変更するクエリ: a[p] から i を削除して、a[c] に i を挿入する
  • 文字 c が区間 [left, right) に含まれるかどうか: a[c] の lower_bound(left) と、lower_bound(right) とを比較する

という風にしてできる。

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

const int INF = 1<<30;
int N, Q;
string S;
vector<set<int> > pls;

void solve() {
    pls.assign(26, set<int>({INF}));
    for (int i = 0; i < N; ++i) pls[S[i] - 'a'].insert(i);

    while (Q--) {
        int type; cin >> type;
        if (type == 1) {
            int ind; char nc; cin >> ind >> nc; --ind;
            auto pc = S[ind];
            S[ind] = nc;
            pls[pc - 'a'].erase(ind);
            pls[nc - 'a'].insert(ind);
        }
        else {
            int l, r; cin >> l >> r; --l;
            int res = 0;
            for (int c = 0; c < 26; ++c) {
                if (pls[c].empty()) continue;
                int lid = *pls[c].lower_bound(l);
                int rid = *pls[c].lower_bound(r);
                if (rid > lid) ++res;
            }
            cout << res << endl;
        }
    }
}

int main() {
    cin >> N >> S >> Q;
    solve();
}

解法 (2): BIT を 26 本もつ

set を用いた解法と思想は同じ。set を用いた方法では「削除」「挿入」「区間に含まれる個数」に答えられればよかった。これらはいずれも BIT を用いてもできる。

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

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

const int INF = 1<<30;
int N, Q;
string S;

void solve() {
    vector<BIT<int>> pls(26, BIT<int>(N+10));
    for (int i = 0; i < N; ++i) pls[S[i]-'a'].add(i+1, 1);
    while (Q--) {
        int type; cin >> type;
        if (type == 1) {
            int ind; char nc; cin >> ind >> nc; --ind;
            auto pc = S[ind];
            S[ind] = nc;
            pls[pc - 'a'].add(ind+1, -1);
            pls[nc - 'a'].add(ind+1, 1);
        }
        else {
            int l, r; cin >> l >> r; --l;
            int res = 0;
            for (int c = 0; c < 26; ++c) if (pls[c].sum(l+1, r+1)) ++res;
            cout << res << endl;
        }
    }
}

int main() {
    cin >> N >> S >> Q;
    solve();
}