ABC423-D を Python で解いた

ABC423 D - Long Waiting
https://atcoder.jp/contests/abc423/tasks/abc423_d

非常に厄介な問題で、自分がなぜ解けなかったのかを解明するのにとても時間がかかりました。ChatGPTと二人三脚です。たくさんの収穫がありました。

コンテスト中の動き

C問題までに30分以上かかり、しかもペナルティ3回。残り70分近くをD問題に費やしましたが解けませんでした。パフォーマンスは茶色、レートも再び大幅ダウンです。

D問題の自分のコードを見ていくと致命的な間違いがあちこちにあり、これはどう転んでもコンテスト中の正解は不可能だったと思いました。この前参加したABC421のD問題と全く同じです。

公式解説の通りに実装する

https://atcoder.jp/contests/abc423/editorial/13871
まずはAtCoder公式が用意してくださった解説を読み、その通りに実装しました。T[i] とすべきところを A[i] にしてしまうなどミスもありましたが割とすんなりと AC になりました。

from sortedcontainers import SortedList
N, K = map(int, input().split())
A, B, C = [0 for _ in range(N+1)], [0 for _ in range(N+1)], [0 for _ in range(N+1)]
standby = SortedList()
for i in range(1, N+1):
  a, b, c = map(int, input().split())
  A[i], B[i], C[i]= a, b, c

T = [0 for _ in range(N+1)]
empty = K # 現在の空席数
for i in range(1, N+1):
  # 団体客 i を入店させる
  AA = max(T[i-1], A[i]) # 入店予定時刻
  if empty >= C[i]: # 空席が十分あれば入店時刻が決定する
    T[i] = AA
    empty -= C[i]
    standby.add([T[i] + B[i], C[i], i])
  else: # 空席が足りない場合は退店処理を行う
    while (empty < C[i]):
      top = standby.pop(0)
      time = top[0]
      empty += top[1]
    T[i] = max(AA, time)
    empty -= C[i]
    standby.add([T[i] + B[i], C[i], i])

for t in T[1:]:
  print(t)

この解法をざっくりと解説するとこんな感じでしょうか。まずこの2つの配列を用意します。

• 入店イベントだけを並べた配列
• 退店イベントだけを並べた配列

とりあえず退店の方は放っておいて入店の方だけを見ていきます。入れない客が出てきたら退店の方を見て処理します。

言われてみればその通りで単純な考え方ですが、ではこれをどうやって思いつけばいいのでしょうか。言われて理解できるのとコンテスト中に思いつけるのには雲泥の差があります。

思いつき方を考える

時系列順に全てのイベントを並べるとこんな感じになります。クエリを処理していくタイプの問題で、途中でクエリが増えるイメージでしょうか。

[入店], [入店], [退店], [入店], [退店], [入店], .......

入店イベントと退店イベントは相互に関連しています。

• あるグループの退店時刻はそのグループの入店時刻によって決まる。
• あるグループの入店時刻は A[i]、直前のグループの入店時刻、別のグループの退店時刻によって決まる。

複雑に関連しているので、それぞれのイベントを別々の配列で管理することを考えます。この発想は前回記事にした ABC421-D に似ていると思いました。あのときも僕はまず各イベントの情報を dict 型で持ち、発生時刻を key, 他の情報を value にして一括管理していました。しかしその後、以下の記事で書いたように2つのポインタを別々に動かすことにより2つの配列を独立して扱う発想を知ってそちらを使いました。

あのときは2つ並行して進めていく感じでしたが、今回はまた別です。入店イベントと退店イベントの関係を見てみると「入店イベント→退店イベント」はシンプルな法則で決められる一方「退店イベント→入店イベント」はかなり複雑な法則で決まります。ということから、入店イベントを中心に見ていって困ったときに退店イベントを参照するという公式解説の流れになります。こんな風に考察していけばいいのかなと思いました。ポイントはこんな感じでしょうか。

• 2つのイベントを一括して扱うのではなく2つの配列に分けて扱う
• ポインタの進め方を考える。今回は一方に着目する。どちらに着目すれば楽になるか比較して考える。

自分の解法

次に、僕が本番中にやろうとしていた方針のままで解くことができないか考えてみました。クエリとして考えて時系列順に処理していきます。

• 入店できるならすぐに入店させる
• 入店できないなら待ち行列に並べる
• 退店イベントが来たらそのときに待ち行列から入店させる

公式解説が入店イベント主導の考え方（入店できないときだけ退店イベントを見る）なら、こちらは退店イベント主導の考え方（退店イベントが来るたびに入店イベントの処理を行う）といえるのかもしれません。が、退店を中心に見る一方で入店イベントも順次処理していきます。こうして見てみるとかなり複雑なことをやってしまっていますね。実際僕はコンテスト中、待ち行列に人が並んでいるときにも一つめの「入店できるならすぐに入店する」を実行してしまっていました。バグが起こりやすい、よくない解法ですね。公式の方がずっと簡潔です。

ここからはこの解法の改善を図ります。もう少し深く考えてみると、「すぐに入店させる」と「待ち行列から入店させる」を区別するからややこしくなることに思い当たります。

• 入店イベントが来たら必ず待ち行列に並べる
• 退店イベントが来たら店の空き人数を増やす
• どちらのイベントが来たときにも、その後に待ち行列をチェックする。入れるなら入店させる。

とする。こうすれば幾分スッキリしますね。「即入店」が例外になってしまっていたので共通化しました。この方針で実装していきました。

SortedList が重い

WA は出なかったので方針はおそらく合っていたのですが、TLE が発生しました。ここは ChatGPT に相談しました。何往復もやりとりした結果、SortedList が重いから heapq を使うようにと言われました。今まで特に深く考えておらず便利な SortedList ばかり使ってきたのですが、どうやらこれは多機能な代わりに動作が重いようです。

• SortedList は挿入・削除・インデックスアクセスができて便利だが、定数倍が重い
• heapq は挿入と最小値の取り出ししかできないが、高速に動作する

恥ずかしいことに全然知りませんでした。全部 SortedList で済ませていました。確かに今回の問題では最も早い時刻のイベントを取り出せればいいので heapq で十分ですね。こんな簡単なことすら知らなかったようです。

実装

というわけで以上のことを盛り込んで実装です。こうして見るとそこまで複雑でもないですね。公式解説のような2つの配列を扱うやり方にしなくても今回は一応なんとかなったことがわかります。条件分岐を整理してシンプルにする工夫と、heapq について知っておくことが必要でした。

余談ですが最初は待ち行列 standby も SortedList で作っていました。これはそもそも Sorted である必要がない（追加順は必ず時刻順になる）から deque でいいよと指摘され、そのように直しました。確かにそうですね。

import heapq
from collections import deque
N, K = map(int, input().split())
A, B, C = [0 for _ in range(N+1)], [0 for _ in range(N+1)], [0 for _ in range(N+1)]
event = []
for i in range(1, N+1):
  a, b, c = map(int, input().split())
  A[i], B[i], C[i] = a, b, c
  heapq.heappush(event, [a, b, c, i])
  
standby = deque([])
person = 0 # 現在の客数
ans = [0 for _ in range(N+1)]
time = 0 # 現在の時刻
while (event):
  [time, end, num, idx] = heapq.heappop(event)
  if num > 0: # 入店イベントが来たとき
    standby.append(idx)
  elif num < 0: # 退店イベントが来たとき
    person += num

  # 待ち行列のチェック
  while (standby and person + C[standby[0]] <= K):
    # 入店イベントを消化し、退店イベントを作って入れる
    top = standby.popleft()
    person += C[top]
    ans[top] = time
    heapq.heappush(event, [time + B[top], 0, -C[top], top])

for an in ans[1:]:
  print(an)

まとめ

○○法といった典型的なアルゴリズム、数学、算数のようなパズルを解く力を要求するような問題ではなく、効率よく処理を進めるためのデータ構造選びや条件分岐・ループの組み方といったまさにプログラミングの基礎力が問われる問題だったと思います。僕の中ではこういうのこそがもっともプログラミングコンテストっぽい問題だと勝手に思っており、こういう問題こそ解けるようになっておきたいです。実際、仕事でプログラミングをしようと思ったらこういう問題が解けることは非常に重要な気がします。

今回に限らず、頭の中である程度ロジックは組めているのに実装ができずにもどかしい思いをすることは多々あるので克服したいところですね。