BEN2のブログ

たまに書いています。

ABC 157 E - Simple String Queries

ABC 157 E - Simple String Queries についてのメモを残します。

セグメント木とは何なのか

この記事で、セグ木について、典型的な問題 Range Minimum Query (RMQ) の解説とともに丁寧に説明されています。自分もこの記事で勉強しました。ありがとうございます。

このあと説明する ABC 157 E - Simple String Queries の実装は RMQ によく似ています。

セグ木 + ビットフラグ

英小文字の種類は 26 しかないので、ビットフラグで「どの英小文字が登場しているか」という情報を保持することにする。
たとえば、次のビットフラグの場合、

00000000000000000000001001

0 番目と 3 番目のビットが立っており、a と d が登場していることを表す。

まずはセグ木の前処理について。
与えられた文字列が abcdbbd のとき、セグ木の葉には、下図のようにビットフラグを保持させる。

f:id:BEN2suzuka:20200302155522p:plain

各節点は、二つの子の論理和を保持するとよい。

f:id:BEN2suzuka:20200302161840p:plain

結局、文字列が与えられた時点では下図のようになる。

f:id:BEN2suzuka:20200302162030p:plain

次にクエリ処理について。
たとえば、区間  [ 2, \, 5) に含まれる英小文字の種類数を答えるときは、下図の橙色の  [ 2, \, 4) [ 4, \, 5)論理和をとり、立っているビットの本数を答えればよい。

f:id:BEN2suzuka:20200302213026p:plain

また、文字列  S i 文字目を書きかえるときは、まずは葉を書きかえる。

f:id:BEN2suzuka:20200302162547p:plain

そして、その親が二つの子の論理和をとり、さらにその親が ... と根まで書きかえていく。

f:id:BEN2suzuka:20200302162807p:plain

RMQ の実装も セグメントツリー入門 で説明されています。

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

int num = 1;
vector<int> A;

// S[i] を c に書きかえる
void update(int i, char c) {
  i += num - 1;
  A.at(i) = 1 << (c - 'a');  // ビットフラグ
  while (i > 0) {
    i = (i - 1) / 2;
    A.at(i) = A.at(i*2 + 1) | A.at(i*2 + 2);
  }
}

// 「区間 [a, b) でどの文字が登場するか」を保持したビットフラグを返す
int query(int a, int b, int k, int l, int r) {
  if (r <= a || b <= l) return 0;
  if (a <= l && r <= b) return A.at(k);
  else {
    int c1 = query(a, b, 2*k+1, l, (l+r)/2);
    int c2 = query(a, b, 2*k+2, (l+r)/2, r);
    return c1 | c2;
  }
}

int main() {
  int N; cin >> N;
  string S; cin >> S;
  num = 1;
  while (num < N) num *= 2;
  A = vector<int>(2*num - 1, 0);  // 完全二分木
  for (int i = 0; i < N; i++) update(i, S.at(i));

  int Q; cin >> Q;
  for (int i = 0; i < Q; i++) {
    int com; cin >> com;
    if (com == 1) {
      int idx; cin >> idx; idx--;
      char c; cin >> c;
      update(idx, c);
    }
    if (com == 2) {
      int left, right; cin >> left >> right;
      left--;
      int bit = query(left, right, 0, 0, num);
      int cnt = 0;
      for (int i = 0; i < 26; i++) {
        if ((bit >> i) & 1) cnt++;  // i 番目のビットが立っている
      }
      cout << cnt << endl;
    }
  }
}

セグ木は不要

YouTube 解説放送 で説明されています。
C++ の標準ライブラリにある set は (平衡) 二分探索木で、値の追加・削除・探索は  O ( \log N ) の計算量で可能。

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

int main() {
  int N; cin >> N;
  string S; cin >> S;
  // 各アルファベット (a, b, c, ...) が S の何文字目に出現するかを保持
  vector<set<int>> sets(26);
  for (int i = 0; i < N; i++) {
    int alp = S.at(i) - 'a';
    sets.at(alp).insert(i);
  }
  int Q; cin >> Q;
  for (int i = 0; i < Q; i++) {
    int com; cin >> com;
    if (com == 1) {
      int idx; cin >> idx; idx--;
      char c; cin >> c;
      sets.at(S.at(idx) - 'a').erase(idx);
      sets.at(c - 'a').insert(idx);
      S.at(idx) = c;
    }
    if (com == 2) {
      int left, right; cin >> left >> right;
      left--; right--;
      int cnt = 0;
      for (int i = 0; i < 26; i++) {
        // 各アルファベットが区間 [left, right] 内で出現するか
        auto iter = sets.at(i).lower_bound(left);
        if (iter != sets.at(i).end() && *iter <= right) cnt++;
      }
      cout << cnt << endl;
    }
  }
}