【競技プログラミング】ダブリングまとめ

競技プログラミングにおける典型アルゴリズムの1つ「ダブリング」の概要と類題を自分用にまとめる。

ダブリング

アルゴリズムの概要

全体の要素数が$NN$個あって、それぞれの要素について、その要素から1回遷移(移動)したときの移動先が定まっているとする。
このとき、「$KK$回遷移したときの到達点」を高速に求めるアルゴリズムの1つがダブリングである。

このような問題を愚直に求めようとすると$O(K)\mathcal{O}(K)$になるが、ダブリングを用いることにより$O(Nlog⁡K)\mathcal{O}(N\log{K})$で計算できるようになる。

アルゴリズムの詳細

一般的なダブリングでは、動的計画法で用いるDPテーブルに似た以下のようなテーブルを事前に計算する。

$dp[i][j]dp[i][j]$: $jj$番目の要素から$2i2^i$回遷移したときの到達地点。ただし、$1≤j≤N,0≤i≤⌈log⁡2K⌉1 \leq j \leq N, 0 \leq i \leq \lceil \log_{2} {K} \rceil$。

このテーブルの初期条件は以下のようになる。ここで、集合$S={1,2,…,N}S = \{1, 2, \ldots, N\}$に対して、「$jj$番目の要素から1回遷移したときの到達地点」を表す関数$f:S→Sf: S \rightarrow S$を$f(j)f(j)$ $(1≤j≤N)(1 \leq j \leq N)$と定める。

$${初期条件: dp[0][j]=f(j)(1≤j≤N)漸化式: dp[i][j]=dp[i−1][dp[i−1][j]](0≤i≤⌈log⁡2K⌉,1≤j≤N) \begin{cases} \text{初期条件: } &dp[0][j] = f(j) \quad (1 \leq j \leq N) \\ \text{漸化式: } &dp[i][j] = dp[i - 1][dp[i - 1][j]] \quad (0 \leq i \leq \lceil \log_2 K \rceil, 1 \leq j \leq N) \end{cases}$$

初期条件と漸化式について補足

初期条件$(i=0)(i = 0)$については、$20=12 ^ 0 = 1$回遷移したときの到達点は$f(j)f(j)$の値をそのまま入れておけばよいというだけである。

漸化式の方については、以下のように導出すればよい。

全ての$jj$と$i−1i - 1$以下について$dp[i−1][j]dp[i - 1][j]$が求まっている状態で、$jj$番目の要素から$2i2 ^ i$回遷移したときの到達点を求めたい。
これを求めるには、$jj$番目の要素から$2i−12 ^ {i - 1}$回遷移した後、更にもう$2i−12 ^ {i - 1}$回遷移した結果が得られればよい。
「$jj$番目の要素から$2i−12 ^ {i - 1}$回遷移したときの到達点」は$dp[i−1][j]dp[i - 1][j]$に入っているので、そこから更に$2i−12^{i - 1}$回遷移した結果は$dp[i−1][dp[i−1][j]]dp[i - 1][dp[i - 1][j]]$に格納されている。これを$dp[i][j]dp[i][j]$に格納すればよい。

事前計算によってこのテーブルを計算できれば、$KK$を2進数とみなし、$KK$の各ビットを下位から見ていき、テーブルを用いて結果を更新していけばよい。
すなわち、$K=2c1+2c2+2c3+⋯+2ck(0≤c1≤c2⋯ck)K = 2^{c_1} + 2^{c_2} + 2^{c_3} + \cdots + 2^{c_k} \quad (0 \leq c_1 \leq c_2 \cdots c_k)$と表したとき、$c1c_1$から順に、答え$AA$を$A=dp[cl][A](1≤l≤k)A = dp[c_l][A] \quad (1 \leq l \leq k)$と更新していくことで解を得られる。

計算量は、事前計算で$O(Nlog⁡K)\mathcal{O}(N\log{K})$、解の計算で$O(log⁡K)\mathcal{O}(\log{K})$となる。

類題

ダブリングを使って解くことのできる類題をいくつか挙げる。

AtCoder Beginner Contest 167 D - Teleporter

問題

AtCoderのページより引用する。

問題文

高橋王国には$NN$個の町があります。町は$11$から$NN$まで番号が振られています。

それぞれの町にはテレポーターが$11$台ずつ設置されています。町$i(1≤i≤N)i (1 \leq i \leq N)$のテレポーターの転送先は町$AiA_i$です。

高橋王は正の整数$KK$が好きです。わがままな高橋王は、町$11$から出発してテレポーターをちょうど$KK$回使うと、どの町に到着するかが知りたいです。

高橋王のために、これを求めるプログラムを作成してください。

制約

• $2≤N≤2×1052 \leq N \leq 2 \times 10^5$
• $1≤Ai≤N1 \leq A_i \leq N$
• $1≤K≤10181 \leq K \leq 10^{18}$

解答例

ダブリングを使って簡単に解くことができる。

テーブルを以下のように定義する。

$dp[i][j]dp[i][j]$: 町$jj$から$2i2^i$回テレポートして到達した町の番号 $(0≤i≤log⁡2K,1≤j≤N)(0 \leq i \leq \log_2 K, 1 \leq j \leq N)$

初期条件と漸化式は以下のようになる。

$${初期条件: dp[0][j]=Aj(1≤j≤N)漸化式: dp[i][j]=dp[i−1][dp[i−1][j]](1≤i≤⌈log⁡2K⌉,1≤j≤N) \begin{cases} \text{初期条件: } &dp[0][j] = A_{j} \quad (1 \leq j \leq N) \\ \text{漸化式: } &dp[i][j] = dp[i - 1][dp[i - 1][j]] \quad (1 \leq i \leq \lceil \log_2 K \rceil, 1 \leq j \leq N) \end{cases}$$

Pythonによるコード例を以下に示す。尚、$KK$が最大で$101810^{18}$であるため$log⁡K\log K$を60程度に取らないといけないが、$NN$は最大で$2×1052 \times 10^5$であり、計算量が$10710^7$オーダーになる場合がある。そのため、PythonではPyPy3で提出しないとTLEになるので注意。

d.py

from typing import List


def main():
    # 入力受け取り
    N, K = map(int, input().split())
    # 0-indexedで受け取っておく
    A: List[int] = list(map(lambda x: int(x) - 1, input().split()))

    # ダブリングのテーブル
    dp: List[List[int]] = [[0 for j in range(N)] for i in range(61)]
    # 初期条件
    for j in range(N):
        dp[0][j] = A[j]

    # 遷移
    for i in range(1, 61):
        for j in range(N):
            dp[i][j] = dp[i - 1][dp[i - 1][j]]

    # 解を求める
    answer: int = 0
    # 現在見ているビットの下位からの桁
    i: int = 0
    # Kを2進数とみなして計算する
    while K:
        # Kの下位からi桁目が1なら遷移する
        if K & 1:
            answer = dp[i][answer]
        # 1つビットシフトする
        K >>= 1
        # iを進める
        i += 1

    # 解の出力(1-indexedに戻す)
    print(answer + 1)


if __name__ == "__main__":
    main()

実際の提出はこちら。

競プロ典型90問 058 - Original Calculator (★4)

問題

AtCoderのページより引用する。

問題文

あなたは奇妙な電卓を持っています。この電卓は$00$以上$105−110^5-1$以下の整数を$11$つ表示できます。この電卓にはボタンAと呼ばれるボタンがあります。整数$xx$が表示されているときに ボタンA を$11$回押すと、次の処理が順番に行われます。

1. $xx$を十進法で表したときの各桁の和を計算し、$yy$とする。
2. $x+yx+y$を$10510^5$で割ったあまりを計算し、$zz$とする。
3. 表示されている整数を$zz$に変更する。

例えば、$9999999999$が表示されているときに ボタンA を$11$回押すと、$99999+(9+9+9+9+9)=10004499999+(9+9+9+9+9)=100044$なので、表示される整数は$4444$に変更されます。

今、この電卓に$NN$が表示されています。 ボタンA を$KK$回押した後に表示されている整数を求めて下さい。

制約

• $0≤N≤105−10 \leq N \leq 10^5-1$
• $1≤K≤10181 \leq K \leq 10^{18}$
• 入力はすべて整数

解答例

ダブリングテーブルを以下のように定義する。

$dp[i][j]dp[i][j]$: 整数$jj$が表示されている状態から$2i2^i$回ボタンを押した後に表示される整数$(0≤i≤log⁡2K,0≤j<105)(0 \leq i \leq \log_2 K, 0 \leq j < 10^5)$

整数$xx$が表示されているときにボタンを1回押した後表示される数値$zz$を求める関数を$f(x)f(x)$と表記することにする。
このとき、初期条件と漸化式は以下のようになる。

$${初期条件: dp[0][j]=f(j)(0≤j<105)漸化式: dp[i][j]=dp[i−1][dp[i−1][j]](1≤i≤⌈log⁡2K⌉,0≤j<105) \begin{cases} \text{初期条件: } &dp[0][j] = f(j) \quad (0 \leq j < 10^5) \\ \text{漸化式: } &dp[i][j] = dp[i - 1][dp[i - 1][j]] \quad (1 \leq i \leq \lceil \log_2 K \rceil, 0 \leq j < 10^5) \end{cases}$$

Pythonによる実装例を以下に示す。こちらはPyPy3でなくてもTLEにはならない様子。

058.py

from typing import List
import math


def f(x: int) -> int:
    # 整数xを文字列として解釈して計算する
    y: int = sum(map(lambda s: ord(s) - ord("0"), str(x)))
    return (x + y) % 100000


def main():
    # 入力受け取り
    N, K = map(int, input().split())

    # log_2 KをMとおく。
    M: int = math.ceil(math.log2(K))

    # ダブリングテーブル
    dp: List[List[int]] = [[0 for j in range(100000)] for i in range(M + 1)]

    # 初期条件
    for j in range(100000):
        dp[0][j] = f(j)

    # 遷移
    for i in range(1, M + 1):
        for j in range(100000):
            dp[i][j] = dp[i - 1][dp[i - 1][j]]

    # 解を求める
    answer: int = N
    i: int = 0
    while K:
        if K & 1:
            answer = dp[i][answer]
        K >>= 1
        i += 1

    print(answer)


if __name__ == "__main__":
    main()

実際の提出はこちら。

AtCoder Beginner Contest 241 E - Putting Candies

問題

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

問題文

長さ$NN$の数列$A=(A0,A1,…,AN−1)A=(A_0,A_1,\ldots,A_{N-1})$が与えられます。
最初の時点では空の皿があり、高橋君は次の操作を$KK$回繰り返します。

• 皿の中のアメの個数を$XX$とする。皿に$A(X mod N)A_{(X\bmod N)}$個のアメを追加する。
  ただし、$X mod NX\bmod N$で$XX$を$NN$で割った余りを表す。

$KK$回の操作の後で、皿の中には何個のアメがあるか求めてください。

制約

• $2≤N≤2×1052 \leq N \leq 2\times 10^5$
• $1≤K≤10121 \leq K \leq 10^{12}$
• $1≤Ai≤1061 \leq A_i\leq 10^6$
• 入力はすべて整数である。

解答例

ダブリングテーブルを以下のように定義する。

$dp[i][j]dp[i][j]$: 皿の上のアメの数が$j(=Xmod N)j (= X \mod N)$個のとき、そこから$2i2^i$回操作を繰り返すことで追加されるアメの総数 $(0≤j≤N,0≤i≤⌈log⁡2K⌉)(0 \leq j \leq N, 0 \leq i \leq \lceil \log_2 K \rceil)$

アメの数が$j(=Xmod N)j (=X \mod N)$のとき、その状態から$2i2^i$回操作を行ったときに追加されるアメの総数は、「皿の上のアメの数が$jj$個のときから$2i−12^{i - 1}$回操作を繰り返したときに追加されるアメの個数」$++$「そこから更にもう$2i−12^{i - 1}$回操作を繰り返したときに追加されるアメの個数」である。
後者の項の解釈：皿の上には元々$jj$個乗っていて、その状態で$2i−12^{i - 1}$回操作を繰り返す(前者の項の操作)と、皿の上には$dp[i−1][j]dp[i - 1][j]$個のアメが追加される。そこから更に$2i−12^{i - 1}$回操作を繰り返すとき、皿の上には$j+dp[i−1][j]j + dp[i - 1][j]$個のアメがあることになるので、$dp[i−1][(j+dp[i−1][j])mod N]dp[i - 1][(j + dp[i - 1][j]) \mod N]$となる。

従って、初期条件と漸化式は以下のようになる。

$${初期条件: dp[0][j]=Aj(1≤j≤N)漸化式: dp[i][j]=dp[i−1][j]+dp[i−1][dp[i−1][j]](1≤i≤⌈log⁡2K⌉,0≤j≤N) \begin{cases} \text{初期条件: } &dp[0][j] = A_j \quad (1 \leq j \leq N) \\ \text{漸化式: } &dp[i][j] = dp[i - 1][j] + dp[i - 1][dp[i - 1][j]] \quad (1 \leq i \leq \lceil \log_2 K \rceil, 0 \leq j \leq N) \end{cases}$$

Pythonによる実装例を以下に示す。この問題も$KK$の最大値が$101210^12$なので、$ii$の値はせいぜい40まででよい。そうしないと全体の計算量が$10710^7$オーダーになり、Pythonだと時間制限が厳しくなるので注意。

e.py

from typing import List


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

    # ダブリングテーブル
    dp: List[List[int]] = [[0 for j in range(N)] for i in range(41)]
    # 初期条件
    for j in range(N):
        dp[0][j] = A[j]

    # 遷移
    for i in range(1, 41):
        for j in range(N):
            dp[i][j] = dp[i - 1][j] + dp[i - 1][(j + dp[i - 1][j]) % N]

    # 答えを格納する変数
    answer: int = 0
    i: int = 0
    while K:
        if K & 1:
            # dp[i][j]は「操作によって追加されるアメの数」であるので、
            # 答えの数に加算していく。
            answer += dp[i][answer % N]
        K >>= 1
        i += 1

    # 回答出力
    print(answer)


if __name__ == "__main__":
    main()

実際の提出はこちら。

AtCoder Beginner Contest 258 E - Packing Potatoes

問題

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

問題文

ベルトコンベアに載って$1010010^{100}$個のじゃがいもが$11$個ずつ流れてきます。流れてくるじゃがいもの重さは長さ$NN$の数列$W=(W0,…,WN−1)W = (W_0, \dots, W_{N-1})$で表され、$i (1≤i≤10100)i \, (1 \leq i \leq 10^{100})$番目に流れてくるじゃがいもの重さは$W(i−1) mod NW_{(i-1) \bmod N}$です。ここで、$(i−1) mod N(i-1) \bmod N$は$i−1i - 1$を$NN$で割った余りを表します。

高橋君は、まず空の箱を用意し、次のルールに従ってじゃがいもを順番に箱に詰めていきます。

• じゃがいもを箱に入れる。箱に入っているじゃがいもの重さの総和が$XX$以上になったら、その箱には蓋をし、新たに空の箱を用意する。

$QQ$個のクエリが与えられます。$i (1≤i≤Q)i \, (1 \leq i \leq Q)$番目のクエリでは、正整数$KiK_i$が与えられるので、$KiK_i$番目に蓋をされた箱に入っているじゃがいもの個数を求めてください。問題の制約下で、蓋をされた箱が$KiK_i$個以上存在することが証明できます。

制約

• $1≤N,Q≤2×1051 \leq N, Q \leq 2 \times 10^5$
• $1≤X≤1091 \leq X \leq 10^9$
• $1≤Wi≤109 (0≤i≤N−1)1 \leq W_i \leq 10^9 \, (0 \leq i \leq N - 1)$
• $1≤Ki≤1012 (1≤i≤Q)1 \leq K_i \leq 10^{12} \, (1 \leq i \leq Q)$
• 入力は全て整数

解答例

公式解説に倣い、各$i=1,2,…,10100i = 1, 2, \ldots, 10^{100}$に対して$j=(i−1)mod Nj = (i - 1) \mod N$をそのじゃがいもの「種類」と呼ぶこととする。
また、配列$C=(C1,C2,…,CN)C = (C_1, C_2, \ldots, C_N)$の要素$CjC_j$を「種類$jj$のじゃがいもから始めて順番に箱詰めを行ったとき、何個入れた時点で蓋をされるか？」を表す数とする。

ダブリングテーブルを以下のように定義する。

$dp[i][j]dp[i][j]$: 種類$jj$のじゃがいもから始めて順番に箱詰めを行っていき、$2i2^{i}$個の箱の蓋が閉じられたときの、次に箱詰めする対象となるじゃがいもの種類。$(0≤i≤⌈log⁡2K⌉,1≤j≤N)(0 \leq i \leq \lceil \log_2 K \rceil, 1 \leq j \leq N)$

$KK$が与えられたとき、ダブリングテーブルを使うことで、$K−1K - 1$個の箱の蓋が閉じられた後最初に入れるじゃがいもの種類がわかることになる($KK$番目の箱に最初に入れるじゃがいもの種類がわかるということである)。
配列$CC$が既知であるなら、これらの情報より、$KK$番目の箱に入るじゃがいもの個数を求めることができる。

配列$CC$の求め方

まずは配列$CC$を求める。これがわからないとダブリングテーブルの漸化式が求まらない。

愚直に求めるならば、各$jj$からスタートして$jj$を1つずつインクリメントしながら$W(j−1)mod NW_{(j - 1) \mod N}$を加算していき、初めて重さの和が$XX$以上になったところでストップする。
このときの添字を$ll$としたとき、$Cj=l−j+1C_j = l - j + 1$と計算することができる。

しかし、この計算だと例えば$X=109,Wj=1X = 10^9, W_j = 1$のような条件では$10910^9$オーダーの計算が必要になってしまう。
そこで、次のように考える。$W=(W1,W2,…,WN)W = (W_1, W_2, \ldots, W_N)$の総和を$SS$とおく。種類$jj$のじゃがいもから始めて順番に箱詰めしていったとき、$S<XS < X$が成り立つならば、1周してきて種類$jj$に戻ってからまたいくつかのじゃがいもを入れて蓋をすることになる。
これを考えると、種類$jj$から始めて蓋をするまでに、$⌊XS⌋\lfloor \frac{X}{S} \rfloor$周してからもう数個じゃがいもを入れることになる。
$⌊XS⌋\lfloor \frac{X}{S} \rfloor$周したとき、残りの許容重量は$Xmod SX \mod S$となるから、種類$jj$から始めて重さが初めて$Xmod SX \mod S$以上になったときの種類を$ll$とすればよいことがわかる。

これを踏まえると、配列$CC$は次のように求められる。

• $CC$の各要素は$⌊XS⌋×N\lfloor \frac{X}{S} \rfloor \times N$で初期化しておく。
• $WW$を2周繋げた配列の累積和を取った配列$VV$を用意しておき、各$jj$について以下のように$ll$を求める。
  • $Vl≥X+VjV_l \geq X + V_j$を満たす最小の添字$ll$を二分探索法によって見つける。
  • $CjC_j$に$(l−j)(l - j)$を加算する。

これで配列$CiC_i$を求めることができた。

ダブリングテーブルの初期条件と漸化式

ダブリングテーブルの初期条件と漸化式は以下のようになる。

$${初期条件: dp[0][j]=(j+Cj)mod N(1≤j≤N)漸化式: dp[i][j]=dp[i−1][dp[i−1][j]](1≤i≤⌈log⁡2K⌉,1≤j≤N) \begin{cases} \text{初期条件: } &dp[0][j] = (j + C_j) \mod N \quad (1 \leq j \leq N) \\ \text{漸化式: } &dp[i][j] = dp[i - 1][dp[i - 1][j]] \quad (1 \leq i \leq \lceil \log_2 K \rceil, 1 \leq j \leq N) \end{cases}$$

あとは$K−1K - 1$回遷移したときの添字$ll$を求め、$ClC_l$を出力すればよい。

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

e.py

from typing import List
import sys
import itertools
import bisect


input = sys.stdin.readline


def main():
    # 入力受け取り
    N, Q, X = map(int, input().split())
    W: List[int] = list(map(int, input().split()))
    K: List[int] = [int(input()) for i in range(Q)]

    # Wの総和
    S: int = sum(W)
    # 配列C
    # 各要素は「何周するか×1周の個数」で初期化しておく
    C: List[int] = [(X // S) * N] * N

    # ダブリングテーブル
    dp: List[List[int]] = [[0 for j in range(N)] for i in range(41)]

    # 周回分はすでに加算しているので残った余りについて計算する
    X %= S
    # Wを2周分繋げて累積和を取った配列V
    V: List[int] = [0] + list(itertools.accumulate(W * 2))
    # 二分探索で配列Cの各要素の値を求める
    for j in range(N):
        l: int = bisect.bisect_left(V, X + V[j])
        C[j] += l - j

    # aダブリングテーブルの初期条件
    for j in range(N):
        dp[0][j] = (j + C[j]) % N

    # 遷移
    for i in range(1, 41):
        for j in range(N):
            dp[i][j] = dp[i - 1][dp[i - 1][j]]

    # 各クエリの解を求める
    for k in K:
        # K - 1番目の箱の蓋を締め終わったとき、
        # 次に入れるじゃがいもの種類が知りたい
        k -= 1
        index: int = 0
        i: int = 0
        while k:
            if k & 1:
                index = dp[i][index]
            k >>= 1
            i += 1

        # 解の出力
        print(C[index])


if __name__ == "__main__":
    main()

実際の提出はこちら。