小猿的击鼓传花

$K(K\ge3)$ 个猿辅导的老师们在玩一个击鼓传花的小游戏。每击一次鼓，拿着花的老师要将花交给别人，不能留在自己手中。游戏开始前花在小猿手中，求击了 $N$ 次鼓后，这朵花又回到小猿手中的方案数，请输出这个数模1000000007后的结果。

解析：在击鼓 $N$ 次后，花要么在小猿手上，要么不在。设在小猿手上的方案数是 $a_n$ ，不在小猿手上的方案数是 $b_n$ ，递推关系是：
第 $n+1$ 次，花只能从其余 $K-1$ 个人的手里传到小猿手中，因此 $a_{n+1}=b_n$ 。如果花仍然不在小猿手上，要么是从小猿手上传来的，要么是从另外 $K-2$ 个人手上传来的，因此 $b_{n+1}=(K-2)b_n+(K-1)a_n$ 。联系可得： $a_{n+2}=(K-2)a_{n+1}+(K-1)a_n$ 。这是一个二阶差分方程，可以通过特征根方程求解。 $\lambda^2=(K-2)\lambda+(K-1)$
如果用上述关系式直接求解，时间复杂度为 $O(N)$ ，只能通过70%的例子。

import java.util.Scanner;

public class Main{

    public static void main(String[] args){
        Scanner input;
        int N, K;

        input = new Scanner(System.in);
        while(input.hasNext()){
            N = input.nextInt();
            K = input.nextInt();
            System.out.println(Solution(N, K));
        }
    }

    private static long Solution(int N, int K){
        int i, M;
        long ans, last, llast;

        M = 1000000007;
        if(N == 1)
            return 0;
        if(N == 2)
            return K - 1;
        i = 2;
        llast = 0;
        last = K - 1;
        ans = 0;
        while(i < N){
            ans = ((K - 2) * last + (K - 1) * llast) % M;
            llast = last;
            last = ans;
            i++;
        }
        return ans;
    }
}

在一般的网站，如果时间复杂度达到 $10^9$ 的数量级，基本都会超时，因此要对时间复杂度进行优化。如果用特征根方程进行求解，对大多数K来说特征根都是无理数，不利于求解。还有另一种解法类似于高阶差分方程的解法，用矩阵来计算。 $\\∵\qquad a_{n+2}=(K-2)a_{n+1}+(K-1)a_n\\ ∴\qquad [a_{n+2},\;a_{n+1}]=[a_{n+1},\;a_{n}]A\\ 其中A=\left[ \begin{matrix}K-2 & 1 \\K-1 & 0 \\\end{matrix} \right]\\可得:\qquad [a_{n},\;a_{n-1}]=[a_{2},\;a_{1}]A^{n-2}=[a_{2},\;a_{1}]\left[ \begin{matrix}A_{11} & A_{12} \\A_{21} & A_{22} \\\end{matrix} \right]\\a_n=a_2\times A_{11}+a_1\times A_{21}=(K-1)\times A_{11}$
所有可以通过求矩阵的幂来计算，这个算法的时间复杂度可以压缩到 $O(logN)$ 。
考虑到递归算法比较耗资源，在这里用迭代算法求矩阵的幂。

import java.util.*;

public class Main{

    public static void main(String[] args){
        Scanner input;
        int N, K;

        input = new Scanner(System.in);
        while(input.hasNext()){
            N = input.nextInt();
            K = input.nextInt();
            System.out.println(new Main().Solution(N, K));
        }
    }

    private long Solution(int N, int K){
        int M;
        long[][] A, An;

        M = 1000000007;
        if(N == 1)
            return 0;
        if(N == 2)
            return K - 1;
        A = new long[][]{{K - 2, 1}, {K - 1, 0}};
        An = pow(A, N - 2);
        return An[0][0] * (K - 1) % M;
    }

    private long[][] pow(long[][] A, int n){
        Stack<Integer> stack;
        long[][] ans;

        stack = new Stack<>();
        ans = new long[][]{{1, 0}, {0, 1}};
        while(n > 0){
            stack.push(n % 2);
            n /= 2;
        }
        while(!stack.isEmpty()){
            multi(ans, ans);
            if(stack.pop() == 1)
                multi(ans, A);
        }
        return ans;
    }

    private void multi(long[][] A, long[][] B){
        long a00, a01, a10, a11;
        int M;

        M = 1000000007;
        a00 = (A[0][0] * B[0][0] + A[0][1] * B[1][0]) % M;
        a01 = (A[0][0] * B[0][1] + A[0][1] * B[1][1]) % M;
        a10 = (A[1][0] * B[0][0] + A[1][1] * B[1][0]) % M;
        a11 = (A[1][0] * B[0][1] + A[1][1] * B[1][1]) % M;
        A[0][0] = a00;
        A[0][1] = a01;
        A[1][0] = a10;
        A[1][1] = a11;
    }
}