猫でもわかるWebプログラミングと副業

本業エンジニアリングマネージャー。副業Webエンジニア。Web開発のヒントや、副業、日常生活のことを書きます。

AOJ 1327: One-Dimensional Cellular Automaton

問題

一次元のセルオートマトンがある。セルは0からN-1までのN個ある。 それぞれのセルには0以上でMより小さい整数の値を持つ。 その値は時間によってかわり、i番目のセルの時間tの値は以下の式によって決まる。 A,B,Cは定数であり、非負の整数である。 i < 0 または N <= i では S(i,t) = 0 とする。 時間Tにおけるオートマトンの値を求めよ。

方針

n = 5 の時を例として考える。以下のような行列を考える。*1

B C 0 0 0
A B C 0 0
0 A B C 0
0 0 A B C
0 0 0 A B

この行列をRとし、オートマトンの状態を行列S(t) で表すと、

S(t) = R * S(t-i)

で表すことができるので、

S(t) = Rm * S(0)

となる。この Rm を普通に求めると mが最大 109 なので、繰り返し二乗法を使い高速化する。これによって O(100 * 100 * log_2(109) となり、log_2(109) はせいぜい30程度なので間に合う。

ソースコード

import java.io.*;
import java.util.*;

class Main {
    static Scanner sc = new Scanner(new InputStreamReader(System.in));
    public static void main(String[] args) throws Exception {
        while(true) {
            // 入力 ----------------------------
            int n = sc.nextInt();
            int m = sc.nextInt();
            int a = sc.nextInt();
            int b = sc.nextInt();
            int c = sc.nextInt();
            int t = sc.nextInt();
            if(n+m+a+b+c+t == 0) break;

            int[] s = new int[n];
            for (int i = 0; i < n; i++) {
                s[i] = sc.nextInt();
            }
            // ここまで入力 -----------------------

            // 行列
            int[][] matrix = new int[n][n];
            for (int i = 0; i < n; i++) {
                matrix[i][i] = b;
            }
            for (int i = 1; i < n; i++) {
                matrix[i][i-1] = a;
                matrix[i-1][i] = c;
            }

            // 配列^t を求める
            matrix = modPow(matrix, t, m);
            
            // 答えとなる行列を求めて出力
            StringBuffer buf = new StringBuffer();
            for (int i = 0; i < n; i++) {
                int ans = 0;
                for (int j = 0; j < n; j++) {
                    ans += matrix[i][j] * s[j];
                }
                buf.append((ans % m) + " ");
            }
            System.out.println(buf.toString().trim());
        }
    }

    /**
    * 繰り返し二乗法を使った行列のpow
    */
    static int[][] modPow(int[][] matrix, int t, int m) {
        matrixMod(matrix, m);
        if(t == 0) {
            int n = matrix.length;
            int[][] ans = new int[n][n];
            for (int i = 0; i < n; i++) {
                ans[i][i] = 1;
            }
            return ans;
        }
        if(t == 1) return matrix;

        if(t % 2 == 0) {
            int[][] mult = matrixMod(matrixMult(matrix, matrix), m);
            return matrixMod(modPow(mult, t/2, m), m);
        } else {
            int[][] mult = matrixMod(matrixMult(matrix, matrix), m);
            return matrixMod(matrixMult(modPow(mult, t/2, m), matrix), m);
        }
    }

    /**
    * aの各要素についてmodを取る非破壊的メソッド
    */
    static int[][] matrixMod(int[][] a, int m) {
        int n = a.length;
        int[][] b = new int[n][n];
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n; j++) {
                b[i][j] = a[i][j] % m;
            }
        }
        return b;
    }

    /**
    * 配列の a * b を行う非破壊的メソッド
    */
    static int[][] matrixMult(int[][] a, int[][] b){
        int n = a.length;
        int[][] c = new int[n][n];
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n; j++) {
                for (int k = 0; k < n; k++) {
                    c[i][j] += a[i][k] * b[k][j];
                }
            }
        }
        return c;
    }
}

*1:tex記法がうまく使えなかったのでコードブロックで行列を表す