ABC391-D を Python で解いた

ABC-391 D - Gravity
https://atcoder.jp/contests/abc391/tasks/abc391_d

先日のD問題が解けなくて悔しかったので翌日またやりました。なんとかACできました。反省の意味も込めて記事にしました。解説というより自分の思考の過程をそのままだらだらと書き連ねているので読みづらいかもしれません。

コンテスト中の自分の動き

Dを見てすぐにはわかりそうになかったのでEも見ました。Eを少し考えてわからなかったのでまたDを見ました。DとEの往復を繰り返し、結局どちらも解けずに3完で低いパフォーマンスに終わってしまいました。残念ですが仕方ありません。反省するのみです。

まずシミュレーションすることを考えました。時間ごとに全てのマスを管理し、揃ったらブロックを消して全ての降り積もったブロックの高さを -1 し、マスは列ごとに配列を作って下から何段目にブロックがあるかを管理し、どの列の最下段にブロックが積もっているのかを管理し……。時間が10^9と大きいので各列の一番下にあるブロックの一番高いやつの高さを見てその分スキップすれば短縮できるというところまで考えてその先が思いつきませんでした。

次にブロックごとに「いつ落ちるのが止まるか」を調べることを考えました。結局あとの質問では実質的にそれを聞かれるわけですから。止まる時刻とそのときの高さをブロックごとに記録します。それとは別に横一列が消える時刻を調べて……と、結局これも断念しました。結局横一列の状態を調べるために毎回 W 回の計算が発生するので無理じゃないかとずっと悩んでいました。

考察

縦も横も大きすぎるのでデータをどうやって管理すればいいのかに悩みましたが、ブロックに着目してみたら見通しが立ちました。縦と横がどんなに広くてもブロックは高々 N 個です。特に時間の経過については座標圧縮のようにしてブロックが下から 1 行目に落ちる時刻ごとに T[0], T[1], T[2], ...... , T[N] みたいにして管理することを考えれば N 回の計算で大丈夫です。

では N 個のブロックをどのようなデータ構造で管理すればいいかですが、x, y 以外に「その列の中で下から何番目」という情報を持たせます。ブロックは各列ごとに初期配置で下にあったものから順番に落ちて降り積もっていきます。この順序は決して変わりません。さらにそのブロックが一番下の行に落ちてきたら「各列で下から p 番目のブロックが一番下の行に落ちたのはこれで計何個か」を p ごとに記録しておきます。これで横一列が揃ったかどうかを高速で判定できるようになります。後で Q 個の質問に答えるためにブロックの番号も欲しいですね。

各種データを管理するための変数

ブロックを管理する B[n][x][y][p] という配列を用意します。また、「各列で下から p 番目のブロックが一番下の行に落ちたのはこれで計何個か」を持っておくために配列 line_up[p] を用意します（「横並び」みたいな意味で調べましたが、よい命名のやり方が思いつきませんでした）。

時刻を管理するために変数 t を用意します。t が変わるたびに配列 B の中身を書き換えていては計算量オーバーになるので、各ブロックの現在の位置は初期値 y から t を引くなどして把握するようにします（これが思いつかなかったためにコンテスト中には解ける見通しが全く立ちませんでした）。というより、実装してみて気づきましたが各ブロックについて「いつ一番下に行き着くのか、落ちるのが止まるのか」がわかればいいので現在位置を把握する必要はありません。

最後に Q 個の質問に答えるために「どのブロックがいつ消滅するのか」を記録する配列がほしいです。ここは block_deleted_time[p] とし、「各列で下から p 番目のブロックが消える時刻」を記録します。p 番目のブロック同士は皆同時に消滅するのでこれがわかれば十分です。

実装

N, W = map(int, input().split())
B = [[-1, -1, -1, -1] for _ in range(N+1)]
for i in range(1, N+1):
  x, y = map(int, input().split())
  B[i] = [i, x, y, 0]
B.sort(key=lambda x : x[2]) # 初期位置の高さの低い順に並び替える
number_from_bottom = [0 for _ in range(W+1)] # 各列で下から何個目まで見たのかを記録する
for i in range(1, N+1): # 初期の高さが低い順に見ていく
  number_from_bottom[B[i][1]] += 1
  B[i][3] = number_from_bottom[B[i][1]]

# 各ブロックが一番下まで落ちきる時刻ごとに現在の状態を見ていく。
# 先の並び替えにより、配列 B は早く落ちるブロックが先に来るように並びかえられている。
# （ここでいう一番下というのは一番下の行もしくは降り積もったブロックのてっぺんに落ちるという意味。
# これまでに横 1 列のブロックの消滅が何度起きたのかは注意しておく必要がある。）
deleted_block_time = [-1 for _ in range(N+1)] # 各列で下から p 番目に落ちてきたブロックが消える時間
deleted_line_num = 0 # これまでに何行分のブロックが消されたか
line_up = [0 for _ in range(N+1)] # 各列で下から p 番目に落ちてきたブロックが一番下まで落ちた個数
for i in range(1, N+1):
  [n, x, y, p] = B[i]
  # ブロックが一番下に落ちる時間 = (y - 1) - [降り積もったブロックの数（床が何段底上げされているか）]
  # つまり (y - 1) - [同じ列の中で自分よりも先に何個のブロックが落ちていったか] + [これまでに消えた行数]
  # なので、
  t = y - 1 - (p - 1) + deleted_line_num
  # 今、この時刻になったものとする。この時点での最下段の状況を見る。
  line_up[p] += 1
  if line_up[p] == W: # 各列で下から p 番目に落ちてきていたブロックが横 1 列に揃ったので消える
    deleted_block_time[p] = t+1
    deleted_line_num += 1

# Q 個の質問に答えるため、何番のブロックがいつ消えるのかを計算する
ans = [-1 for _ in range(N+1)]
for i in range(1, N+1): 
  n, p = B[i][0], B[i][3]
  ans[n] = deleted_block_time[p]

Q = int(input())
for _ in range(Q):
  t, a = map(int, input().split())
  if ans[a] == -1:
    print("Yes")
  else:
    if ans[a] < t+0.5:
      print("No")
    else:
      print("Yes")

おわりに

翌日に挑戦して解けたのですが1時間以上かかりました。コンテスト中に悩んだり、終わってから他人のコメントや解説を読んだおかげでようやく解けたので僕がコンテスト中にこれを解くのは無理だったと思います。緑diff だったので解きたかったのですが厳しいですね。