B - シャンパンタワー / Champagne Tower 解説 by admin
gemini-3.5-flash-thinking概要
この問題は、直列に並んだ \(N\) 個のグラスからなるシャンパンタワーに対して、上からシャンパンを注ぐシミュレーションを行う問題です。グラスが満杯になると、溢れた分が即座に下のグラスへと流れ落ちます。
愚直にシミュレーションを行うと、すでに満杯になったグラスを何度も走査することになり実行時間制限に間に合いません。そこで、Union-Find(素集合データ構造)を利用して「次に空き容量があるグラス」を高速に見つけることで、効率的にシミュレーションを行います。
考察
素朴なアプローチ(シミュレーション)とその限界
シャンパンを注ぐたびに、1段目のグラスから順に「満杯か? 満杯なら次の段へ」という処理を繰り返す素朴な方法を考えます。
例えば、すでに \(1\) 段目から \(N-1\) 段目までのグラスがすべて満杯であるとします。この状態で \(1\) 段目に少量のシャンパンを注ぐと、シャンパンは \(1\) 段目から \(N-1\) 段目までをすべて素通りし、最終的に \(N\) 段目に到達します。 この判定のために毎回 \(N\) 回のステップを踏むことになります。
このような操作を \(Q\) 回繰り返すと、最悪の場合の計算量は \(O(NQ)\) となります。 本問題では \(N, Q \le 2 \times 10^5\) であるため、最悪ケースでは \(4 \times 10^{10}\) 回程度の演算が必要になり、実行時間制限(TLE)になってしまいます。
高速化のアイデア:「満杯のグラスをスキップする」
高速化のために、「すでに満杯になったグラスは、次から探索時にスキップする」というアプローチを考えます。
「グラス \(i\) にシャンパンが流れ込んできたとき、グラス \(i\) 以降で最初にまだ満杯になっていない(空き容量がある)グラスはどこか?」を高速に求めることができれば、無駄な走査を省くことができます。
これは、Union-Find(経路圧縮あり)を用いることで、ほぼ定数時間で求めることができます。
具体的には、各グラス \(i\) に対して、以下のようにグループを管理します。
- parent[i] を「グラス \(i\) 以降で、まだ満杯になっていない最小のグラスのインデックス(代表元)」とする。
- 初期状態では、すべてのグラスに空き容量があるため、parent[i] = i です。
- グラス \(i\) が満杯になった瞬間、グラス \(i\) の親を「グラス \(i+1\) のグループの代表元」にマージします。
これにより、find(i) を呼び出すだけで、グラス \(i\) 以降で次にシャンパンを注ぐべきグラスを一瞬で見つけることができます。
アルゴリズム
1. データの準備
- グラスの容量を格納する配列 \(C\)
- 各グラスの現在のシャンパン量を格納する配列
amount(初期値はすべて0) - Union-Find用の配列
parent- 1-indexedで管理するため、サイズは \(N+2\) とします。
- \(N+1\) 番目の要素は、すべてのグラスが満杯になってタワーからシャンパンが溢れ出たときのための「番兵(受け皿)」として機能します。
2. 操作 \(1\)(注ぐ)の処理
\(1\) 段目のグラスに \(V\) ミリリットルのシャンパンを注ぎます。
curr = find(1)によって、現在一番上にある「満杯でないグラス」を取得します。- \(V > 0\) かつ
curr <= N(タワーの範囲内)である限り、以下の処理を繰り返します。- グラス
currの空き容量space = C[curr] - amount[curr]を計算します。 - ケースA:\(V \ge space\)(グラスが満杯になる場合)
- グラス
currのシャンパン量を満杯(C[curr])にします。 - 注ぐ量を \(V \leftarrow V - space\) に更新します。
- グラス
currは満杯になったので、その親を隣のグラスのグループにマージします:parent[curr] = find(curr + 1)。 - 次に注ぐ先を
curr = find(curr + 1)に更新します。
- グラス
- ケースB:\(V < space\)(グラスが満杯にならない場合)
- グラス
currに残りすべてのシャンパンを注ぎます:amount[curr] += V。 - 注ぐ量がなくなったので \(V = 0\) とし、ループを終了します。
- グラス
- グラス
3. 操作 \(2\)(問い合わせ)の処理
- \(k\) 段目のグラスに現在入っているシャンパンの量
amount[k]をそのまま出力します。
計算量
時間計算量: \(O((N + Q) \alpha(N))\)
- 各グラスが「満杯になる」というイベントは、タワー全体で高々 \(1\) 回しか発生しません。
- したがって、Union-Find のマージ(親の更新)が発生する回数は、すべてのクエリを通して最大でも \(N\) 回です。
- Union-Find の
find操作は、経路圧縮を行うことでアッカーマン関数の逆関数 \(\alpha\) を用いて \(O(\alpha(N))\)(ほぼ定数時間)で行えます。 - よって、全体の時間計算量は \(O((N + Q) \alpha(N))\) となり、実行時間制限に余裕で間に合います。
空間計算量: \(O(N)\)
- 容量配列 \(C\)、現在の量配列
amount、Union-Find用のparent配列はいずれもサイズ \(O(N)\) です。 - したがって、空間計算量は \(O(N)\) となります。
実装のポイント
番兵(\(N+1\) 段目)の導入: 最下段である \(N\) 段目のグラスが満杯になったとき、
find(N + 1)を指すようにマージします。これにより、配列の範囲外参照(Index Error)を防ぐと同時に、「溢れたシャンパンがどこにも溜まらずに失われる」という仕様を綺麗にシミュレートできます。Pythonでの高速入出力: Pythonは入出力のオーバーヘッドが大きいため、
sys.stdin.read().split()で入力を一括で取得し、出力もsys.stdout.writeを使ってまとめて出力することで、実行時間を大幅に短縮しています。ソースコード
import sys
def solve():
# 入力をすべて取得
input_data = sys.stdin.read().split()
if not input_data:
return
N = int(input_data[0])
Q = int(input_data[1])
# 1-indexedにするため、サイズは N + 2 とする
C = [0] * (N + 2)
for i in range(1, N + 1):
C[i] = int(input_data[1 + i])
# Union-Find(Disjoint Set Union)の初期化
# parent[i] は、グラス i 以降でまだ満杯になっていない(空き容量がある)最小のグラスのインデックスを指す
parent = list(range(N + 2))
amount = [0] * (N + 2)
def find(i):
curr = i
while parent[curr] != curr:
curr = parent[curr]
# 経路圧縮
p = i
while p != curr:
nxt = parent[p]
parent[p] = curr
p = nxt
return curr
idx = N + 2
out = []
len_data = len(input_data)
while idx < len_data:
type_q = int(input_data[idx])
if type_q == 1:
V = int(input_data[idx + 1])
idx += 2
curr = find(1)
while curr <= N and V > 0:
space = C[curr] - amount[curr]
if V >= space:
# グラス curr が満杯になる
amount[curr] = C[curr]
V -= space
# 満杯になったので、次の空いているグラスの代表元にマージする
nxt = find(curr + 1)
parent[curr] = nxt
curr = nxt
else:
# グラス curr は満杯にならない
amount[curr] += V
V = 0
else:
k = int(input_data[idx + 1])
idx += 2
out.append(str(amount[k]))
# まとめて出力
sys.stdout.write("\n".join(out) + "\n")
if __name__ == "__main__":
solve()
この解説は gemini-3.5-flash-thinking によって生成されました。
投稿日時:
最終更新: