AtCoder Beginner Contest 167 D問題

AtCoder Python

D問題

難易度：茶色

atcoder.jp


茶色の中でも難し目の問題を解いてみましたが、
TLEとWAのダブルパンチを喰らい、結局以下のサイトを参考にする形となりました。

note.com

K(≤10**5)個の街に設置されているワープマシンを
N(≤10**18)回使用したらどこの街に飛ばされるかという問題です。

Nに関しては(≤10**18)なのでまともに解いていてはTLEです。
Kに関しては(≤10**5)なのでO(N)回の計算量程度は実行可能です。

まず考えるべきなのは、「ループ」です。
途中から同じ街を何度もループし始めます。
この規則性を見破り、ループ+剰余の式を作成し、結果的にK回目のループ先を導きます。

規則性の見つけ方：TLE

正直考え方は良い線でしたが、規則性・ループをまとめ方にTLEを出してしまいました。

itertoolsのCountersを使用したりしても単体の計算量がO(N)の場合、
for文内に埋め込んでしまうと、全体でO(N**2)となってしまいますねえ・・・。

  • 訪れた回数配列：visit
  • 訪れた街を列挙した配列：goal
  • 2回目訪れた街を列挙(ループしてる)：loop

Countersを使用しなくとも、
配列のindexをkey-1、要素をvaluesに対応づけているので、
全体でO(N)で動作します。

特異なケース：WA

剰余 = 序盤の規則性がないゾーン = len(goal) - len(loop)
あとこいつを考慮してませんでした。
WAの原因はこいつです。

K と 剰余の差分から、規則性のある数列に置き換えます。
それにはK > 剰余が前提です。

以下の場合分けを考えないといけませんでしたね。

  • goalがKより大きいケース
  • goalがKより小さいケース

与えられた

from collections import Counter
N, K = map(int,input().split())
tele = list(map(int, input().split()))

#A1から順にワープした時の経由順
now = 0
goal,loop,visit = [],[],[0]*N

while visit[now]!=2:

  if visit[now]==0:
    goal.append(now)
  else:
    loop.append(now)
    
  visit[now] += 1
  now = tele[now]-1

if len(goal)>K:
  print(goal[K]+1)
else:
  amari = len(goal)-len(loop)
  print(loop[(K-amari)%len(loop)]+1)