🍜

ABC268 C - Chinese Restaurant Python解答例

2022/09/11に公開約3,000字

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問題までは全探索が視野に入るっていう通説が以前あった気がするんですが、とうとうそれが通用しなくなって来つつあることを実感しています。

GitHubで編集を提案

Discussion

ログインするとコメントできます