ABC268 C - Chinese Restaurant Python解答例

AtCoder Beginner Contest 268 C - Chinese RestaurantをPythonで解きます。

問題

問題文をAtCoderのページより引用します。

問題文

回転テーブルの周りに人$0$、人$1$、$\ldots$、人$N-1$がこの順番で反時計回りに等間隔で並んでいます。また、人$i$の目の前には料理$p_i$が置かれています。
あなたは次の操作を$0$回以上何度でも行うことが出来ます。

• 回転テーブルを反時計回りに$1$周の$\frac{1}{N}$だけ回す。これによって、(この操作の直前に)人$i$の目の前にあった料理は人$(i+1) \bmod N$の目の前に移動する。

操作を完了させた後において、人$i$は料理$i$が人$(i-1) \bmod N$、人$i$、人$(i+1) \bmod N$のいずれかの目の前に置かれていると喜びます。
喜ぶ人数の最大値を求めてください。

制約

• $3 \leq N \leq 2 \times 10^5$
• $0 \leq p_i \leq N-1$
• $i \neq j$ならば$p_i \neq p_j$
• 入力はすべて整数

料理$p_i$を人$i$の元に持っていくには何回操作が必要か？

$N$の数が多いので、愚直に全探索するとTLEします。なんとかして高速化することを考えます。

説明の便宜上、人$i$が喜ぶときの条件を以下のように言い換えます。

• 人$i$は料理$i$が人$(i - 1) \mod N$の目の前に置かれていると喜ぶ$\rightarrow$料理$i$が人$i$の左隣にあるとき喜ぶ
• 人$i$は料理$i$が人$i$の目の前に置かれていると喜ぶ$\rightarrow$料理$i$が人$i$の正面にあるとき喜ぶ
• 人$i$は料理$i$が人$(i + 1) \mod N$の目の前に置かれていると喜ぶ$\rightarrow$料理$i$が人$i$の右隣にあるとき喜ぶ

まず、問題の中で定義されている操作の回数は0以上$N$未満であることがわかります。$N$回操作を行うと元の配置に戻るためです。
この範囲の回数だけ操作を行う場合、料理$i$が人$i$の左隣、正面、右隣のいずれかに来るような操作回数はそれぞれ1つしか無いことがわかります。なぜなら、テーブルは反時計回りにしか回らず、1周して戻ってくることも無いからです。

これを踏まえて、人$i$が喜ぶために必要な操作の回数、すなわち「料理$p_i$を人$i$の左隣、正面、右隣に持っていくためには何回の操作が必要か？」を考えてみます。
これはそれぞれ以下のように計算できます。

• 左隣に持っていくための回数：$(N - p_i + i - 1) \mod N$
• 正面に持っていくための回数：$(N - p_i + i) \mod N$
• 右隣に持っていくための回数：$(N - p_i + i + 1) \mod N$

操作回数を求める式の求め方

求める操作回数を$k$とおきます。
料理$p_i$を人$i$の正面に運ぶときの操作を、以下の2ステップに分けます。

1. $p_i$を人0の正面まで運ぶ
2. そこから人$i$の正面まで運ぶ

1.を達成するために必要な操作回数を$k_1$、2.の達成のために必要な回数を$k_2$とおきます。
$k_1$については、$k_1$回操作して$p_i$が人0の位置に戻るため、$p_i + k_1 = N$が成り立ちます。従って、$k_1 = N - p_i$となります。
そこから人$i$の正面に回すにはちょうど$i$回の操作が必要になるので、$k_2 = i$となります。

従って、必要な操作回数は$k = k_1 + k_2 = (N - p_i + i) \mod N$となります。
尚、$N$を法としているので、$(N - p_i + i) \mod N \equiv (p_i - i) \mod N$が成り立ちます(公式解説にかかれている式の形です)。

左隣と右隣については正面についての式にそれぞれ$-1$、$+1$することで求めることができます。

これらの回数だけ操作を行ったときだけ、人$i$は喜ぶことができます。言い換えれば、これらの回数以外の回数だけ操作を行ったとき、人$i$が喜ぶことはありません。

操作を$k$回行ったときの、喜ぶ人数の合計を$C_k$とします。
人$i$に着目すると、上に示した3通りの回数だけ操作を行ったときだけ人$i$が喜ぶので、喜ぶ人数の合計が1増えます。
すなわち、$C_{(N - p_i + i - 1) \mod N}, C_{(N - p_i + i) \mod N}, C_{(N - p_i + i + 1) \mod N}$を1だけ増やせます。

全ての$i$について、対象の$C_k$を増やしていくことで、全ての操作回数における喜ぶ人数の合計を求めることができ、その中の最大値が答えとなります。

実装例

Pythonによる実装例を以下に示します。

c.py

from typing import *


def main():
    # 入力受け取り
    N: int = int(input())
    P: List[int] = list(map(int, input().split()))

    # 操作をk回行ったときの喜ぶ人数の合計を表す配列
    C: List[int] = [0] * N
    # 全ての人iについて計算する
    for i in range(N):
        # 料理p_iが人iの左隣に来るような回数だけ操作を行うと、人iが喜ぶのでCが1増える
        C[(N - P[i] + i - 1) % N] += 1
        # 料理p_iが人iの正面に来るような回数だけ操作を行うと、人iが喜ぶのでCが1増える
        C[(N - P[i] + i) % N] += 1
        # 料理p_iが人iの右隣に来るような回数だけ操作を行うと、人iが喜ぶのでCが1増える
        C[(N - P[i] + i + 1) % N] += 1

    # 最大値を出力する
    print(max(C))


if __name__ == "__main__":
    main()

実際の提出はこちら。

ABCのC問題までは全探索が視野に入るっていう通説が以前あった気がするんですが、とうとうそれが通用しなくなって来つつあることを実感しています。