ケーキ同好会でも、プログラミングは必須スキルです！

ケーキ同好会12

「ちゃんとみんな来たね」

6月なのにもう暑いねー、と駅のホームでひまり先輩が言った。

「はーい、お菓子は何円までですか？」

「なぜ今聞く」

みんなは楽しそうだけど、私はコンテストのことが不安で仕方がない。

やば、緊張で少しおなか痛くなってきたかも。

そんな私の気持ちとは裏腹に、楽し気な接近メロディが鳴り、電車が近づいてくるのが見えた。

大丈夫。できるだけのことはやったはず。

「じゃあ、行こうか」

ホームに到着した電車に乗り込み、空いていたボックス席にみんなで座った。

椅子に座るとき、横に座ったひまり先輩が私の顔を見たような気がした。

そのあと、うーん、と考えるような仕草をして、

「駅といえば、こんな問題はどうかな？」

......なぜか問題を出してきた。

問題文

$N$ 個の駅があります。駅は 1 から $N$ まで番号が振られています。

駅 $i\ (1\leq i \leq N)$ を出発した電車は次に駅$A_i$に到着します。

正の整数$K$が与えられます。次の条件を満たす駅$i$の個数を求めてください。

- 駅 $i$ から出発した電車がちょうど $K$ 回駅を移動して駅 $i$ に戻ってくる。

制約

入力はすべて整数

$2 \leq N \leq 2\times 10^5$

$1 \leq A_i \leq N$

$i \neq A_i$

$1 \leq K \leq 10^{18}$

入力形式

$N\ K$

$A_1\ A_2\ \dots\ A_N$

入力例1

$\texttt{4 6}$

$\texttt{3 3 4 1}$

出力例1

$\texttt{3}$

入力例2

$\texttt{6 437193755483746372}$

$\texttt{6 6 2 5 3 5}$

出力例2

$\texttt{4}$

問題考えてた方が気が楽かも。入力例1を考えてみよう。

【挿絵表示】

駅1から出発して、$K = 6$回駅を移動すると、

\[1 \rightarrow 3 \rightarrow 4 \rightarrow 1 \rightarrow 3 \rightarrow 4 \rightarrow 1\]

となって、駅1にちょうど戻ってくる。

こんな感じで、元の駅に戻ってくる出発駅が何個あるか数えればいいんだね。

「これって、それぞれの駅から$K$回移動して、元の駅に戻ってくるかをfor文で数えたら良さそうですね」

「正しい。けど、それだと入力例2が実行時間に間に合わない」

「そのアルゴリズムの計算量は$O(NK)$で、最大で$2\times10^{23}$くらいの計算が必要になっちゃうよ！」

確かにそうだった。

1秒で$10^8$回くらいの計算回数だったね。

計算量も考慮に入れながらプログラムを書かないと。

じゃあ、素因数分解のときにやったみたいに、ちょっとコードを変更したら改善するかな？

「今回も解法の高速化を考えるけど、条件文一つ変えるくらいではだめで、根本的なアルゴリズム自体の高速化が必要」

「高速化の手法としてはいろいろあるけど、

「何回もやってる処理を減らす」、「計算に必要なデータを前計算する」、「適切なデータ構造を使う」とかをよくやるね」

「その方法は「アルゴリズムとデータ構造」のような教科書に載ってる」

Pythonの本もあるのに、アルゴリズムの本もやらないといけないの？！Pythonの本と同じくらいの分厚さだったら、挫折してしまいそう。

「でも、その方法を適用する方法を考えることも同じくらい大切だね。大きい順に並べ替えたらうまくいくとか、問題自体の性質を使う必要があるよ」

なるほど、じゃあ今考えたアルゴリズムも、無駄な処理を何回もしてるのかな？

うーん、ちょっと考えてみたけど、何回もしてる処理とかはなさそう。

「じゃあ、問題の解法をどうやって高速化していくかの考え方を教えるよ」

うーん、でも今のアルゴリズムで、無駄な処理をしてるところとかあるのかな？

「やっぱりまずは具体例からだね。試しに入力例1のものを考えてみよう！」

駅1から順番に移動すると、

\[1\rightarrow3\rightarrow4\rightarrow1\rightarrow3\rightarrow4\rightarrow1\dots \]

みたいに、同じ場所を何回も通ってることが分かる。

【挿絵表示】

「同じ場所を何回も通るってことは、何回も同じ処理を繰り返してるってことだから、無駄なところだよね」

「ほんとだ！」

「でも、同じ場所を通らないこともあるんじゃない？」

「少なくとも$N+1$回移動すれば、2回通った駅が必ず一つある」

「そうなの？」

「分かった！駅は$N$個しかないからでしょ？」

「正解！」

「移動してるときに同じ駅に来たことが分かれば、また戻ってくるまでの距離はまとめて進めることができる」

「これによって計算量は$O(N\cdot\min(N,K))$になるよ」

\[1\rightarrow 3\rightarrow 4 \rightarrow 1 \xrightarrow{3\times 10\text{回移動}} 1 \rightarrow 3\]

「実際、前の解法だと$10^{23}$程度の計算回数だったのに対し、この解法だと最大で$10^{10}$くらいになって10兆倍くらい高速になった」

「すごい......！」

でも、まだ$10^{10}$だと数分程度かかるから、もっと高速化しないといけないね。最初の解法から10兆倍速くなったのに、まだ速くしないといけないんだ。なかなか難しい問題だね。

「どうしても1回ずつ進むのは効率的ではないので、ある程度長い距離を一度に進みたい」

「どうすればいいと思う？」

「ヒントは、ある駅から次の駅は１つしかないことに注目するよ」

長い距離を一気に進みたいのか......

1回ずつ進むんじゃなくて、何かの倍数だけ進むとか？でも、何かの倍数だけ進んだ場所を求めるには、結局1回ずつ進むしかないからなあ。うーん、分かりません。

ひまり先輩が私の耳元で囁いた。

「2を$10^{12}$回かけるには？」

思ってもみなかった問題が出てきたため、私は驚いて立ち上がった。

「え、繰り返し二乗法！？」

「駅$i$から$2^j$回進んだ駅の場所を先に計算して配列に保存し、$K$回移動するときはその配列を使って移動する」

\[1\xrightarrow{1\text{回移動}}3\xrightarrow{2\text{回移動}}1\xrightarrow{4\text{回移動}}3\]

「これで計算量は$O(N\log K)$になる。計算回数は$10^7$回くらいで、最初のアルゴリズムからは1京倍速くなった」

「速くなりすぎでしょ！？」

n, k = map(int,input().split())

a = list(map(int,input().split()))

logk = 60

dp = [[0]*n for _ in range(logk)]

for i in range(n):

dp[0][i] = a[i]-1

for i in range(1, logk):

for j in range(n):

dp[i][j] = dp[i-1][dp[i-1][j]]

now = list(range(n))

for i in range(logk):

if(k&1):

for j in range(n):

now[j] = dp[i][now[j]]

k >>= 1

res = 0

for i in range(n):

if(now[i] == i):

res += 1

print(res)

●

「着いたー！五国大学！」

駅について数分歩くと、私たちの高校とは違う、大きなキャンパスが広がっていた。受付はどこらへんかな、と見回していると、遠くの建物の入り口に人が立っているのが見えた。

「なんかめっちゃこっちに手を振ってる人いるんですけど」

「うん？」

哀先輩が私の方を振り返った。

「あそこがコンテストの受付だよ」

「え？なんで私たちが参加するって判ったんですか？」

「ひまちゃんを見ればすぐ判る」

うん？後ろを振りかえってみると、ひまり先輩もめちゃくちゃ手を振っていた。

たくさんの棒人間が手をくねくねさせたTシャツを着た男子の方に歩いて行った。近くまで近づいたとき、首から下げた名札に王城と書かれているのが見えた。

「王城さん！お久しぶりです！」

とひまり先輩が元気な挨拶をした。

「ひまり先輩の知りあいですか？」

と哀先輩にこっそり聞いた。

「情報オリンピックのチューターの人」

先輩ってことね。と哀先輩が私に耳打ちした。ひまり先輩の方を見たら、名札をいっぱい持って私たちの方に戻ってきていた。

「はい、みんな名札とってねー」

名札にはケーキ同好会という手書きの文字と、ショートケーキのイラストが書かれていた。

......王城さんが書いたのかな？

「じゃあ、建物の中入ろうか」