AtCoder が提供しているABC（AtCoder Beginner Contest）341 のE問題をC++で解いてみました。ABC341は、2024年2月17日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 Alternating String（Difficulty : 1223）
問題はリンク先をご覧ください。
一工夫して、セグメント木を使います。AtCoder Problems による Difficulty は 1223 でした。
解答案
文字列 $S$ の $i$ 文字目と $i+1$ 文字目が同じであれば値 $0$ を持ち、異なれば値 $1$ を持つ配列を計算します。入力例1をもとに計算してみます。
入力例1 文字列 $S$ の初期値は以下です。
1 0 1 0 0
差分の配列は、以下となります。
1 1 1 0
最初に、クエリ2（1、3）を調べます。
1 0 1 0 0
差分配列の該当する箇所は、以下となります。出力は、Yes です。
1 1 1 0
次のクエリ2（1、5）を調べます。
1 0 1 0 0
差分配列の該当する箇所は、以下となります。出力は、No です。
1 1 1 0
次にクエリ1（1、4）で文字列を書き換えます。
0 1 0 1 0
差分配列を更新します。変化した文字を赤字にします。
1 1 1 1
再びクエリ2（1、5）を調べます。
0 1 0 1 0
差分配列の該当する箇所は、以下となります。出力は、Yes です。
1 1 1 1
次にクエリ1（3、3）で文字列を書き換えます。
0 1 1 1 0
差分配列を更新します。変化した文字を赤字にします。
1 0 0 1
次のクエリ2（2、4）を調べます。
0 1 1 1 0
差分配列の該当する箇所は、以下となります。出力は、No です。
1 0 0 1
差分配列の値を見れば、以下が分かります。
- クエリ1は、文字列を置き換えても、差分の配列は、端の2か所しか変化しない。
- 0と1が反転しても、差分配列の途中の値は変わらない。
- クエリ2は、対応する差分の配列がすべて1であれば Yes を、それ以外の場合は No を出力する。
C++ プログラム例（ABC341E）
セグメント木を使えば、端の値を更新しながら、指定された範囲の合計を求めることができます。この記事で紹介しています。差分配列をセグメント木に載せます。
注意
この問題では、クエリ1とクエリ2の添え字は、1からカウントしています。セグメント木の要素は、[0, N] の領域を確保しておき、差分配列を [1, N-1] = [1, N) に格納すれば、1からカウントする場合と整合が取れます。上に示した例と照らし合わせて確認してください。
- 文字列 $S$ の $i$ 文字目と $i+1$ 文字目の差分配列をセグメント木 $i+1$ 番目の要素に格納します（80ー81行目）。
- クエリ1（L、R）に対して、セグメント木の L-1 番目の要素と R 番目の要素を更新します。注意としては、LとRは問題に与えられた1からカウントした値を使います（89ー90行目）。
- クエリ2（L、R）に対して、セグメント木の [L, R) の範囲の値がすべて1になっているか（つまり R-L になっているか）調べます（92行目）。
以下が、C＋＋プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
class SegmentTree {
private:
vector<int> data;
int (*op)(int, int);
int e;
public:
int size;
SegmentTree(int n, int (*_op)(int, int), int _e) {
op = _op;
e = _e;
size = 1;
while (size < n) {
size *= 2;
}
data.assign(2 * size - 1, e);
}
int get(int k) {
return data[k + size - 1];
}
void update(int k, int a) {
k += size - 1;
data[k] = a;
while (k > 0) {
k = (k - 1) / 2;
data[k] = op(data[k * 2 + 1], data[k * 2 + 2]);
}
}
int query(int a, int b) {
return query_sub(a, b, 0, 0, size);
}
// call ST.query(a, b, 0, 0, size);
int query_sub(int a, int b, int k, int l, int r) {
if ((r <= a)||(b <= l)) {
return e;
}
if ((a <= l)&&(r <= b)) {
return data[k];
} else {
int vl = query_sub(a, b, k * 2 + 1, l, (l + r) / 2);
int vr = query_sub(a, b, k * 2 + 2, (l + r) / 2, r);
return op(vl, vr);
}
}
};
int op(int a, int b)
{
return a + b;
}
int main()
{
int n, q;
cin >> n >> q;
string s;
cin >> s;
if (n == 1) {
for (int i = 0; i < q; ++i) {
int type, L, R;
cin >> type >> L >> R;
if (type == 2) {
cout << "Yes" << endl;
}
}
return 0;
}
SegmentTree st(n + 1, op, 0);
for (int i = 0; i < n - 1; ++i) {
if (s[i] != s[i + 1]) {
st.update(i + 1, 1);
}
}
for (int i = 0; i < q; ++i) {
int type, L, R;
cin >> type >> L >> R;
if (type == 1) {
st.update(L - 1, 1 - st.get(L - 1));
st.update(R, 1 - st.get(R));
} else if (type == 2) {
if (st.query(L, R) == R - L) {
cout << "Yes" << endl;
} else {
cout << "No" << endl;
}
}
}
return 0;
}AtCoder では、AtCoder Library（ACL）と呼ばれるライブラリが提供されています。ACL で、セグメント木クラス segtree も提供されています。以下は、segtree を使ったプログラムです。
#include <bits/stdc++.h>
#include <atcoder/segtree>
using namespace std;
using namespace atcoder;
int op(int a, int b)
{
return a + b;
}
int e(void)
{
return 0;
}
int main()
{
int n, q;
cin >> n >> q;
string s;
cin >> s;
if (n == 1) {
for (int i = 0; i < q; ++i) {
int type, L, R;
cin >> type >> L >> R;
if (type == 2) {
cout << "Yes" << endl;
}
}
return 0;
}
segtree<int, op, e> st(n + 1);
for (int i = 0; i < n - 1; ++i) {
if (s[i] != s[i + 1]) {
st.set(i + 1, 1);
}
}
for (int i = 0; i < q; ++i) {
int type, L, R;
cin >> type >> L >> R;
if (type == 1) {
st.set(L - 1, 1 - st.get(L - 1));
st.set(R, 1 - st.get(R));
} else if (type == 2) {
if (st.prod(L, R) == R - L) {
cout << "Yes" << endl;
} else {
cout << "No" << endl;
}
}
}
return 0;
}どちらも AC（Accepted＝正しいプログラム）と判定されました。
Python のセグメント木について
Python は、セグメント木を標準ライブラリとしては用意していません。このブログの Python プログラムは、なるべく言語標準のライブラリで実装したいと考えています。このような事情で、この記事では、Python に移植したプログラムの紹介を省略します。
最後に
更新するクエリと結果を返すクエリを処理する問題でした。セグメント木を使うことは予想ができましたが、下準備を思いつきませんでした。
先週にセグメント木について記事を書いたため（記事1、記事2、記事3）、応用例を学んで使いこなせるようになりたいです。
引き続き ABC の問題を紹介していきます。
