ダブリング

はじめまして。
株式会社Acompanyの高橋 昂汰です。

この記事は「OPTIMIND x Acompany Advent Calendar 2021」 7日目の記事です。

今日は趣味でやっている競技プログラミングの問題を解いていて面白いなと思った問題を紹介したいと思います。

※問題概要はリンクからご覧ください

D - 阿弥陀

初見で見た時は行列累乗を用いて解けばいいと思いました。

問題の一つ目のサンプルを見ると
行を選んだ線、列を結果で表すと1回目のあみだくじの結果は
$${\begin{pmatrix}0 & 0 & 0 & 1& 0  \\ 0 & 1 & 0 & 0& 0 \\ 0 & 0 & 0 & 0& 1 \\ 0 & 0 & 1 & 0& 0 \\ 1 & 0 & 0 & 0& 0\end{pmatrix}}$$
と表すことができます。
これを用いると$${a_k}$$をk回後の結果とすると、

$${a_k= \begin{pmatrix}1 & 0 & 0 & 0& 0  \\ 0 & 1 & 0 & 0& 0 \\ 0 & 0 & 1 & 0& 0 \\ 0 & 0 & 0 & 1& 0 \\ 0 & 0 & 0 & 0& 1\end{pmatrix} * \begin{pmatrix}0 & 0 & 0 & 1& 0  \\ 0 & 1 & 0 & 0& 0 \\ 0 & 0 & 0 & 0& 1 \\ 0 & 0 & 1 & 0& 0 \\ 1 & 0 & 0 & 0& 0\end{pmatrix}^k}$$

このように式を立てることでk回後の結果が求まります.

しかしこれでは行列積に$${\Omicron (N^3)}$$、繰り返し2乗法を用いても$${logD}$$回計算を行うため全体で$${\Omicron (N^3logD)}$$の計算量となるので到底間に合いません。

解法

この問題の解答としてはダブリングを用いることで高速に求めることができます。
ダブリングとは
1回目の結果を利用して2回目の結果を求める
2回目の結果を利用して4回目の結果を求める
4回目の結果を利用して8回目の結果を求める
8回目の結果を利用して16回目の結果を求める
…
のように$${2^k}$$回目の結果がわかっていれば$${2^{k+1}}$$回目の結果を求められることを利用したアルゴリズムです。
先ほど述べた繰り返し２乗法と同じような考え方です。

今回の問題だと
$${dp[k][i]}$$を$${i}$$からあみだくじを始めて$${2^k}$$回繰り返した時の結果とします。
$${dp[0][i]}$$を1回あみだくじを行った結果を元に初期化しておきます。

$${dp[k][i]}$$から$${dp[k+1][i]}$$への遷移は
$${dp[k+1][i]=dp[k][dp[k][i]]}$$
とすることができる。

これを全ての$${0 \leqq i < n}$$についてk=40まで求めておきます
C++で書くと以下のようなコードになります

for(int k=0;k<40;k++){
    for(int i=0;i<n;i++){
        dp[k+1][i]=dp[k][dp[k][i]];
    }
}

これでDが$${2^k}$$で表すことができる時は求めることができました。

しかしDが$${2^k}$$で表せない時は求めることができていません。
そこでDを二進数表記した時に各桁のビットが立っていれば処理を実行すると考えます。

例としてD=7の時を考えます。
7回繰り返すのは1回行う、2回行う、4回行うと分けて考えることができるので結果は$${dp[0][dp[1][dp[2][i]]]}$$で求まります。
これは7の二進数表記は111=100+10+1と$${2^j}$$の形だけで表せることができるからです。

このようにどのようなDでも$${2^j}$$の形の和だけで表すことができるのでビットが立っている桁だけ処理を行うことで答えを求めることができます。

この処理をC++で書くと以下のようになります。

for(int i=0;i<n;i++){
    int ans=i;
    for(int k=0;k<=40;k++){
        if((d>>k)&1){  // dをkだけ左にシフトして、k桁目のビットが立っていれば処理を行う
            ans=dp[k][ans];
        }
    }
　　　　　　　　　cout << ans << endl;
}

計算量は前処理で$${\Omicron (NlogD)}$$、答えを求めるのに$${\Omicron (NlogD)}$$なので全体として$${\Omicron (NlogD)}$$で解くことができました。

ソースコード

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using ld = long double;
using P  = pair<long long,long long>;
#define rep(i,n) for(ll i=0;i<n;i++)
ll linf=1e18;
int main(){
    ll n,m,d;
    cin >> n >> m >> d;
    
    vector<ll> a(m);
    vector<ll> start(n); 
    for(ll i=0;i<n;i++){
        start[i]=i;
    }
    for(ll i=0;i<m;i++){
        cin >> a[i];
        a[i]--;
        swap(start[a[i]],start[a[i]+1]);
    }
    vector<vector<int> > dp(41,vector<int>(n,-1));
    //初期化
    for(ll i=0;i<n;i++){
        dp[0][start[i]]=i;
    }
    //前処理
    for(ll k=0;k<40;k++){
        for(ll i=0;i<n;i++){
            dp[k+1][i]=dp[k][dp[k][i]];
        }
    }
    
    for(int i=0;i<n;i++){
        int ans=i;
        for(int k=0;k<=40;k++){
            if((d>>k)&1){
                ans=dp[k][ans];
            }
        }
        cout << ans+1<<endl;
    }
}   

まとめ

水色になったもののまだまだ知らないアルゴリズムが多いので今後も青色目指して精進していきたいと思います。