ABC435C：発想を得るためのステップガイド

はじめに

本記事では ABC 435 C 問題を取り扱います。
本問は一見すると単純ですが、

• ドミノが連鎖して倒れるという物理的イメージをコードに落とし込むこと
• どこをどう高速化すべきか発想すること

といった点に難しさを感じる方もいらっしゃることと思います。
そこで本記事では、愚直な解法から問題の骨子を抽出し、高速な解法を導き出す過程を整理して説明します。
なかなか発想が出ない場合の行動指針として参考になれば幸いです。

問題

全文：C - Dominos

高さ $A_i$ のドミノが等間隔に並んでおり、ドミノ $i$ が倒れると区間 $[i, i + A_i)$ にあるドミノがすべて倒れます。このとき、ドミノ $1$ を倒したときに結果として倒れるドミノの数を求める問題です。
各ドミノが別のドミノを倒す区間を持っていることを意識するのが解くためのポイントとなります。

解法に至るための行動と考え方

もし解法がすぐに思いつかない場合、以下の手順を踏むことで問題解決の発想に至りやすくなります。

解法がわからなくてもできること

• 入力例を図に描く
• まずは愚直に実装する
• ランダムテストケースを作成して検証する

考察において重要なこと

• 愚直な実装の計算量を見積る
• 制約から目標の計算量を定める
• 愚直な解法において「無駄な計算」をしている箇所を探す
• その無駄な計算について、必要な記録を保持したり、探索法を工夫したりして高速化することを考える

実践例

以下に各アクションの具体例を示します。

入力例を図解する

まずは与えられた入力例を図に起こしてみて、問題文がどんな特徴を持っているのか把握します。

図 1: 倒れたドミノが後続のドミノを倒す様子

1 番目のドミノが倒れると 2, 3 番目のドミノが倒れ、3 番目のドミノが倒れることで 4 番目のドミノも倒れます。このとき、3 番目のドミノが「自分が倒すドミノの範囲」によって「全体の倒れる範囲」を伸ばしていることに注意します。

愚直な解法の実装例

まずは単純な全探索・シミュレーションのコードを書いていきます。とりあえず愚直に実装してみることで、問題文の読解が深まり、典型問題に落とし込む道筋が見えるかもしれません。
不必要に高速化せず、問題文の通りに実装することが実装のコツです。

def brute_force(n, a):
    is_fell = [False] * n
    is_fell[0] = True
    for i in range(n):
        if not is_fell[i]:
            continue
        for j in range(i, min(i + a[i], n)):
            is_fell[j] = True
    ans = is_fell.count(True)
    return ans

このコードを実行した場合の最悪ケースは $A_i$ が常に最大値 $N$ を取る場合で、二重 for ループの実行回数は $\frac{1}{2} N(N + 1)$、計算量は $\mathcal{O}(N^2)$ となります。
制約より、これを $\mathcal{O}(N)$ まで改善すれば十分高速な解法となります。

ランダムなテストケースを作成

ランダムなテストケースを生成し、愚直な解法と高速な解法との間で結果を比較します。
発想が思い浮かばない間にテストコードを書いておくことで、時間の浪費を防ぎつつ、後で実装する高速な解法の信頼性を高めることが可能です。

def test():
    import random

    # 実行速度に合わせて適当な回数だけ検証する
    for _ in range(10):
        n = 10
        a = [random.randint(1, n) for _ in range(n)]
        # solve1(n, a) は高速な解法とする
        ans = solve1(n, a)
        ans_bf = brute_force(n, a)
        if ans != ans_bf:
            # 愚直な解法と一致しない場合のテストケースは要確認
            # デバッグ・コーナーケースへの対応のヒントになる
            print("WA")
            print(f"{n=}, {a=}: {ans_bf=} but {ans=}")
            break
        else:
            print("AC")

考察

ドミノ $i$ が倒れると、同時に半開区間 $[i, i + A_i)$ のドミノが倒されます。
すなわち、「倒れるドミノ」を処理するたびに、そのドミノが持つ「倒す範囲」が元の範囲に加えられていきます。
したがって、最終的に倒れるドミノの集合は、各ドミノに対応する半開区間の和集合（すべての区間を重ね合わせてカバーされる範囲）によって決まります。

図2: 各ドミノの「倒す範囲」の重ね合わせで「最終的に倒れる範囲」が決まる

愚直な解法においては倒れた各ドミノの区間を全探索していました。しかし実際には倒れることが決まった区間の右端 $\max(r)$ さえ保持していれば十分です。
したがって、倒れることが決まったドミノ $i$ について区間の終点 $r$ の最大値を更新しながらカウントしていけばよいことがわかります。ループ変数 $i$ が倒れることが確定した範囲の右端に到達した場合、それ以降のドミノは倒れません。したがって $i \ge \max(r)$ となった時点で処理を終了すれば求める答えが得られます。
この高速化により、内側のループが解消されて全体で $\mathcal{O}(N)$の計算量となります。

解法例

def solve1(n, a):
    # 1 番目のドミノは必ず倒す
    ans = 1
    # このとき max_r は 1 番目のドミノの区間終点となる
    max_r = a[0]
    for i in range(1, n):
        if i < max_r:
            max_r = max(max_r, i + a[i])
            ans += 1
        else:
            # 倒れるドミノがなくなったら処理を終了する
            break
    return ans

最後にランダムなテストケースで愚直な解法 brute_force(n, a) の結果と比較し、一致することを確認しました。

終わりに

発想に詰まった場合、テストケースの図示、愚直な解法の実装によって問題を具体化し、何を高速化すべきかの方針を立てるのが有効です。
また、ランダムなテストケースを用いて愚直な解法と比較することにより、提出する解法のミスを早期発見し、デバッグのヒントを得られます。結果として問題を解くのにかける時間を短縮することが期待できます。
こうした手順は本問に限らず多くの問題に応用できます。問題にアプローチする足がかりの一つとしてぜひご活用ください。