E - 交互に並べられる区間 / Intervals That Can Be Arranged Alternately Editorial by admin
Gemini 3.0 Flash (Thinking)概要
与えられた白黒の石の並びにおいて、色の個数差が 1 以内であるような連続部分区間(よい区間)を、クエリで与えられた範囲内からいくつ探し出せるかを答える問題です。
考察
1. 「よい区間」の条件を言い換える
「石を自由に並べ替えて、隣り合う色がすべて異なるようにできる」という条件は、石の色の個数に注目するとシンプルになります。 白い石の個数を \(N_W\)、黒い石の個数を \(N_B\) とすると、交互に並べられる条件は \(|N_W - N_B| \leq 1\) であることと同値です。 (例:Wが2個、Bが1個なら「WBW」と並べられますが、Wが3個、Bが1個だとどうしてもWが隣り合ってしまいます)
2. 累積和の利用
各石について、W を \(+1\)、B を \(-1\) と置き換えて考えます。
区間 \([l, r]\) における \(N_W - N_B\) は、この数値の総和です。
累積和配列 \(P\) を \(P[i] = (\text{左から } i \text{ 番目までの数値の和})\) と定義すると、区間 \([l, r]\) が「よい区間」である条件は以下のように書けます。
$\(|P[r] - P[l-1]| \leq 1\)\(
つまり、\)-1 \leq P[r] - P[l-1] \leq 1\(、すなわち **\)P[r]\( と \)P[l-1]\( の差が \)1$ 以内** であればよいことになります。
3. クエリへの対応
各クエリ \((L_i, R_i)\) では、\(L_i \leq l \leq r \leq R_i\) を満たすペア \((l, r)\) のうち、上記の条件を満たすものの個数を求めます。 これは累積和のインデックスで言い換えると、範囲 \([L_i-1, R_i]\) から 2 つのインデックス \(i, j\) (\(i < j\)) を選んだとき、 \(|P[j] - P[i]| \leq 1\) となる組み合わせ数 を求める問題に帰着されます。
素朴に全探索すると各クエリに \(O((R-L)^2)\) かかり、全体で \(O(QN^2)\) となり間に合いません。この「範囲内のペア計数」を効率的に行うために Mo’s algorithm を使用します。
アルゴリズム
Mo’s algorithm
クエリをあらかじめ並び替えて、区間の端点を少しずつ動かしながら答えを更新していくオフラインアルゴリズムです。
- 準備: 累積和 \(P\) を計算します。負の値にならないよう、全体に \(N\) 程度のオフセットを加えておきます。
- クエリのソート: クエリ \([L_i-1, R_i]\) をブロック分割法に基づき、端点の移動距離が最小限になるようソートします。
- 状態の維持:
cnt[v]: 現在の区間に含まれる累積和の値vの出現回数。current_ans: 現在の区間内での条件を満たすペアの総数。
- 伸縮操作:
- 値
valを区間に追加する場合、すでに区間内にあるval-1,val,val+1の個数分だけcurrent_ansが増加します。 - 値
valを区間から除く場合、残っているval-1,val,val+1の個数分だけcurrent_ansが減少します。
- 値
計算量
- 時間計算量: \(O((N + Q)\sqrt{N})\)
- 累積和の構築に \(O(N)\)。
- Mo’s algorithm による端点の移動距離の合計が \(O((N + Q)\sqrt{N})\) です。
- 各移動(追加・削除)は \(O(1)\) で行えます。
- 空間計算量: \(O(N + Q)\)
- 累積和配列、出現回数配列
cnt、クエリの保存に必要です。
- 累積和配列、出現回数配列
実装のポイント
インデックスのオフセット: 累積和 \(P[i]\) は最小で \(-N\) になる可能性があるため、配列の添字として使うために
P[i] += N + 1のように底上げをする必要があります。高速な入出力: \(N, Q\) が \(5 \times 10^4\) と比較的大きいため、C++ の場合は
ios::sync_with_stdio(false); cin.tie(nullptr);を用いて入出力を高速化するのが定石です。Mo’s algorithm の最適化: クエリのソート時に、ブロックごとに右端の移動方向を交互に入れ替える(ジグザグに動かす)ことで、定数倍の高速化が期待できます。
ソースコード
#include <iostream>
#include <vector>
#include <string>
#include <algorithm>
#include <cmath>
using namespace std;
/**
* Problem: Good Intervals in a range [L, R]
* A range is "good" if its stones can be rearranged such that no two adjacent stones have the same color.
* This condition is equivalent to |count(W) - count(B)| <= 1.
* Let P[i] be the prefix sum where 'W' is +1 and 'B' is -1.
* The condition for range [l, r] is |P[r] - P[l-1]| <= 1.
* For each query (Li, Ri), we need to count pairs (i, j) such that Li-1 <= i < j <= Ri and |P[j] - P[i]| <= 1.
* This can be solved efficiently using Mo's algorithm.
*/
struct Query {
int l, r, id, block;
// Standard Mo's algorithm sorting with alternating r direction optimization
bool operator<(const Query& other) const {
if (block != other.block) return block < other.block;
return (block % 2 == 0) ? (r < other.r) : (r > other.r);
}
};
int main() {
// Fast I/O for competitive programming
ios::sync_with_stdio(false);
cin.tie(nullptr);
int N, Q;
if (!(cin >> N >> Q)) return 0;
string S;
cin >> S;
// Calculate prefix sums where 'W' contributes +1 and 'B' contributes -1
vector<int> P(N + 1);
P[0] = 0;
for (int i = 0; i < N; ++i) {
P[i + 1] = P[i] + (S[i] == 'W' ? 1 : -1);
}
// Shift prefix sum values to be non-negative for use as array indices
// Original range is [-N, N], shifted range is [1, 2N+1]
for (int i = 0; i <= N; ++i) {
P[i] += N + 1;
}
// Determine block size for Mo's algorithm
int block_size = max(1, (int)(N / sqrt(Q + 1)));
vector<Query> queries(Q);
for (int i = 0; i < Q; ++i) {
int L, R;
cin >> L >> R;
queries[i].l = L - 1; // Range of indices in P is [Li-1, Ri]
queries[i].r = R;
queries[i].id = i;
queries[i].block = queries[i].l / block_size;
}
// Sort queries based on Mo's algorithm order
sort(queries.begin(), queries.end());
// cnt[v] stores the frequency of value v in the current range of P
// current_ans stores the number of pairs (i, j) in the range satisfying |P[j] - P[i]| <= 1
vector<int> cnt(2 * N + 5, 0);
vector<long long> results(Q);
long long current_ans = 0;
int cur_l = 0, cur_r = -1;
// Add an element to the current range
auto add = [&](int idx) {
int val = P[idx];
// Count existing elements P[k] such that |val - P[k]| <= 1
current_ans += (long long)cnt[val - 1] + cnt[val] + cnt[val + 1];
cnt[val]++;
};
// Remove an element from the current range
auto remove = [&](int idx) {
int val = P[idx];
cnt[val]--;
// Subtract pairs involving the removed element that satisfy the condition
current_ans -= (long long)cnt[val - 1] + cnt[val] + cnt[val + 1];
};
// Process each query using Mo's algorithm pointers
for (int i = 0; i < Q; ++i) {
int L = queries[i].l;
int R = queries[i].r;
while (cur_l > L) add(--cur_l);
while (cur_r < R) add(++cur_r);
while (cur_l < L) remove(cur_l++);
while (cur_r > R) remove(cur_r--);
results[queries[i].id] = current_ans;
}
// Output all results in the order of the original queries
for (int i = 0; i < Q; ++i) {
cout << results[i] << "\n";
}
return 0;
}
この解説は gemini-3-flash-thinking によって生成されました。
posted:
last update: