AtCoder ABC241 E - Putting Candies をフロイドの循環検出やダブリングで解く

https://atcoder.jp/contests/abc241/tasks/abc241_e
をダブリングや循環を検出してときます。

• 方法1: ダブリング
• 方法2: 循環を調べる
• 方法3: フロイドの循環検出アルゴリズム : extra空間計算量O(1)

1: ダブリング

ダブリングするだけです。ダブリングってなんだろう？という場合はSlideShareのAtCoder167Dをダブリングで解くなどを見てください。

ただし、コストを計算する必要があるので、遷移のテーブルとコストのテーブルを持ちます。
遷移のテーブルは先に$nxt_i = a_i mod n$とすると毎回計算しなくて良いです。
costを計算するとき、loopをタブリングテーブルの段数として、今、[loop][i]にいるときは、cost[loop-1][i]に加えて、次の移動先であるnxt[loop-1][i]を踏むコストであるcost[loop - 1][nxt[loop - 1][i]]を足すところが少しごちゃっとするかもしれません。

実装(Python)

n, k = map(int, input().split())
dat = list(map(int, input().split()))
nxt = [[-1] * n for _ in range(61)] # 2^60まで持つ。遷移テーブル
cost = [[-1] * n for _ in range(61)] # 2^60まで持つ。こっちはコストテーブル
for i in range(n): nxt[0][i] = (i + dat[i]) % n  # 1マス先のテーブル
for i in range(n): cost[0][i] = dat[i]  # 1マス先のテーブル
for loop in range(1, 61): # ダブリングテーブルを作る
    for i in range(n):
        nxt[loop][i]= nxt[loop-1][nxt[loop-1][i]]
        cost[loop][i] = cost[loop - 1][nxt[loop - 1][i]] + cost[loop - 1][i]
ans = 0
i = 0
for loop in range(61): # ダブリングをする
    if ((k >> loop) & 1) == 1:
        ans += cost[loop][i]
        i = nxt[loop][i]
print(ans)

2: 循環検出

以下のように考えます。各種を求めた時に、index, カウントがどこまでされたのか？を図に書いて丁寧に考えて実装するといいでしょう。

• 移動を繰り返すと必ずどこかで循環（ループ）に入る。これは訪れたことのあるノードかを管理すれば変わる。
• もし、ループの開始位置とループの長さが分かれば、以下の3つを足せばよい
  • 1: ループに入るまでのコスト
  • 2: ループ1週のコスト $\times$ 何回ループしきれるか？
  • 3: ループしきれない場合、そのループの行けるところまでのコスト

実装(Python)

import sys
n, k = map(int, input().split())
dat = list(map(int, input().split()))
i = 0
cnt = 0
visited = [-1] * n # スタートから、そのノードを踏むまでに最短で何回ループが必要？
while True:
    # 同じノードを踏むまでぐるぐる回す
    if visited[i] != -1: break # 同じノードを見ようとしたときに終わり
    visited[i] = cnt
    cnt += 1
    i = (i + dat[i]) % n
looppoint = i # これが循環するポイント
cycle = cnt - visited[looppoint] # サイクルは、cnt=衝突した回数 - そのポイントの最小ステップ
ans = 0
# まず、ループの入り口まで異動する
i = 0
while i != looppoint:
    if k == 0: # ループの入り口に行く前にkが終わったら終わり
        print(ans)
        sys.exit(0)
    k -= 1
    ans += dat[i]
    i = (i + dat[i]) % n
# ループを一周するコストを計算する
cyclecost = 0
i = looppoint
for _ in range(cycle):
    cyclecost += dat[i]
    i = (i + dat[i]) % n
ans += cyclecost * (k // cycle) # ループをn周した分を加算

# 最後にあまりを足す
k %= cycle
i = looppoint
for _ in range(k):
    ans += dat[i]
    i = (i + dat[i]) % n
print(ans)

3: フロイドの循環検出アルゴリズム

循環検出をみんな大好きフロイドの循環検出アルゴリズムでやります。宛先の書いてある配列を与えると大体$O(N)$で循環の長さ:cycleLenと循環の開始位置:nodeNumLoopStartが求められます。尚、extraな空間計算量はなんと$O(1)$です。つまり、visitedなどの訪れたか？を空間O(N)で持たなくて良いです。

何が嬉しいかというと空間使用量O(1)なので、最小メモリが狙えます。(2022/03/01 00:45時点でACしているコードの中で最小の使用メモリ)

実装(C)

#include <stdio.h>
#include <malloc.h>
#define N	200000
static int a[N];
int n;
int nxt(int cur){
    return (cur + a[cur]) % n;
}

int main() {
    int i;
    long long x, k;
    scanf("%d%lld", &n, &k);
    for (i = 0; i < n; i++) scanf("%d", &a[i]);
    x = 0;
    int hare, tortoise;
    hare = 0;
    tortoise = 0;
    while(1){
        hare = nxt(nxt(hare));
        tortoise = nxt(tortoise);
        if(hare == tortoise) break;
    }
    hare = 0;
    int stepCountStartToLoopStart = 0;
    while(1){
        ++stepCountStartToLoopStart;
        hare = nxt(hare);
        tortoise = nxt(tortoise);
        if(hare == tortoise) break;
    }
    int nodeNumLoopStart = hare;
    int cycleLen = 0;
    hare = nodeNumLoopStart;
    tortoise = nodeNumLoopStart;
    while(1){
        ++cycleLen;
        hare = nxt(nxt(hare));
        tortoise = nxt(tortoise);
        if(hare == tortoise) break;
    }
    long long cycleCost = a[nodeNumLoopStart];
    i = nxt(nodeNumLoopStart);
    while(i != nodeNumLoopStart){
        cycleCost += a[i];
        i = nxt(i);
    }
    i = 0;
    while(i != nodeNumLoopStart && k > 0){
        x += a[i];
        i = nxt(i);
        --k;
    }
    x += cycleCost * (k / cycleLen);
    k %= cycleLen;
    i = nodeNumLoopStart;
    while (k > 0){
        x += a[i];
        i = nxt(i);
        --k;
    }
    printf("%lld\n", x);
    return 0;
}

実装(Python)

def do():
    n, k = map(int, input().split())
    dat = list(map(int, input().split()))
    nxt = [None] * n
    for i in range(n): nxt[i] = (i + dat[i]) % n # next

    # フロイドの循環検出アルゴリズムここから
    hare = tortoise = 0  # 初期位置
    while True:
        hare = nxt[nxt[hare]]  # 2歩進む
        tortoise = nxt[tortoise]  # 1歩進む
        if hare == tortoise: break
    hare = 0  # ウサギだけスタートに戻す
    stepCountStartToLoopStart = 0  # スタートから巡回の開始位置までの距離
    while True:
        stepCountStartToLoopStart += 1
        hare = nxt[hare]  # 1歩進む
        tortoise = nxt[tortoise]  # 1歩進む
        if hare == tortoise: break
    nodeNumLoopStart = hare  # 循環の開始位置
    cycleLen = 0  # 循環の長さ
    hare = tortoise = nodeNumLoopStart
    while True:
        cycleLen += 1
        tortoise = nxt[tortoise]
        hare = nxt[nxt[hare]]
        if tortoise == hare: break
    # フロイドの循環検出アルゴリズムここまで
    # ここから循環１サイクルでかかるコストを計算する
    cycleCost = dat[nodeNumLoopStart]
    i = nxt[nodeNumLoopStart]
    while i != nodeNumLoopStart:
        cycleCost += dat[i]
        i = nxt[i]
    # ここまで循環コスト
    # スタートから循環開始位置まで移動する
    ans = 0
    i = 0
    while i != nodeNumLoopStart:
        ans += dat[i]
        i = nxt[i]
        k -= 1
        if k == 0: # そもそもその前にkが0になればそれが応え
            print(ans)
            return
    # 今、循環のスタート位置にいるので、まず、何周できたかのコスト計算
    ans += cycleCost * (k // cycleLen)
    k %= cycleLen
    # あとはk回回ればいい
    i = nodeNumLoopStart
    while k > 0:
        ans += dat[i]
        i = nxt[i]
        k -= 1
    print(ans)
do()