📚

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

2022/07/11に公開

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

ダブリング

アルゴリズムの概要

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

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

アルゴリズムの詳細

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

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

このテーブルの初期条件は以下のようになる。ここで、集合S = \{1, 2, \ldots, N\}に対して、「j番目の要素から1回遷移したときの到達地点」を表す関数f: S \rightarrow Sf(j) (1 \leq j \leq 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)については、2 ^ 0 = 1回遷移したときの到達点はf(j)の値をそのまま入れておけばよいというだけである。

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

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

事前計算によってこのテーブルを計算できれば、Kを2進数とみなし、Kの各ビットを下位から見ていき、テーブルを用いて結果を更新していけばよい。
すなわち、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)と表したとき、c_1から順に、答えAA = dp[c_l][A] \quad (1 \leq l \leq k)と更新していくことで解を得られる。

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

類題

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

AtCoder Beginner Contest 167 D - Teleporter

問題

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

問題文

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

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

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

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

制約

  • 2 \leq N \leq 2 \times 10^5
  • 1 \leq A_i \leq N
  • 1 \leq K \leq 10^{18}

解答例

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

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

dp[i][j]: 町jから2^i回テレポートして到達した町の番号 (0 \leq i \leq \log_2 K, 1 \leq j \leq 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によるコード例を以下に示す。尚、Kが最大で10^{18}であるため\log Kを60程度に取らないといけないが、Nは最大で2 \times 10^5であり、計算量が10^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のページより引用する。

問題文

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

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

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

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

制約

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

解答例

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

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

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

\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のページより引用する。

問題文

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

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

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

制約

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

解答例

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

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

アメの数がj (=X \mod N)のとき、その状態から2^i回操作を行ったときに追加されるアメの総数は、「皿の上のアメの数がj個のときから2^{i - 1}回操作を繰り返したときに追加されるアメの個数」+「そこから更にもう2^{i - 1}回操作を繰り返したときに追加されるアメの個数」である。
後者の項の解釈:皿の上には元々j個乗っていて、その状態で2^{i - 1}回操作を繰り返す(前者の項の操作)と、皿の上にはdp[i - 1][j]個のアメが追加される。そこから更に2^{i - 1}回操作を繰り返すとき、皿の上にはj + dp[i - 1][j]個のアメがあることになるので、dp[i - 1][(j + dp[i - 1][j]) \mod 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による実装例を以下に示す。この問題もKの最大値が10^12なので、iの値はせいぜい40まででよい。そうしないと全体の計算量が10^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のページより引用する。

問題文

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

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

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

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

制約

  • 1 \leq N, Q \leq 2 \times 10^5
  • 1 \leq X \leq 10^9
  • 1 \leq W_i \leq 10^9 \, (0 \leq i \leq N - 1)
  • 1 \leq K_i \leq 10^{12} \, (1 \leq i \leq Q)
  • 入力は全て整数

解答例

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

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

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

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

配列Cの求め方

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

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

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

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

  • Cの各要素は\lfloor \frac{X}{S} \rfloor \times Nで初期化しておく。
  • Wを2周繋げた配列の累積和を取った配列Vを用意しておき、各jについて以下のようにlを求める。
    • V_l \geq X + V_jを満たす最小の添字lを二分探索法によって見つける。
    • C_j(l - j)を加算する。

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

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

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

\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 - 1回遷移したときの添字lを求め、C_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()

実際の提出はこちら

GitHubで編集を提案

Discussion