HDU 3306 Another kind of Fibonacci

/*
*  [题意]
*   已知:
*       F(0)=1, F(1)=1, F(n)=F(n-1)+F(n-2) (n>=2)
*       A(0)=1, A(1)=1, A(n)=X*A(n-1)+Y*A(n-2) (n>=2)
*   求：S(n), S(n) = (A(0)^2)+(A(1)^2)+...+(A(n)^2)
*  [解题方法]
*   (A(n)^2) = (X*A(n-1) + Y*A(n-2))^2
*            = X*X*(A(n-1)^2) + Y*Y*(A(n-2)^2) + 2*X*Y*(A(n-1)*A(n-2))
*   A(n)*A(n-1) = (X*A(n-1) + Y*A(n-2)) * A(n-1)
*               = X*(A(n-1)^2) + Y*(A(n-1)*A(n-2))
*   S(n) = S(n-1) + A(n)^2
*   所以可以构造矩阵：
*   |1  1   0   0       |       |   S(n-1)      |       |   S(n)    |
*   |0  X*X Y*Y 2*X*Y   |   *   |   (A(n-1)^2)  |   =   |(A(n)^2)   |
*   |0  1   0   0       |       |   (A(n-2)^2)  |       |(A(n-1)^2) |
*   |0  X   0   Y       |       | A(n-1)*A(n-2) |       |A(n)*A(n-1)|
*/
#include <iostream>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <algorithm>
#include <math.h>
using namespace std;
#define FF(i, n) for(int i = 0; i < n; i++)
#define M 4

int ans[M], mod = 10007;
int ret[M][M];
int init[M][M];
int ss[M][M] = {1, 1, 0, 0,
                0, 0, 0, 0,
                0, 1, 0, 0,
                0, 0, 0, 0};

void ini(int n)
{
    FF(i, n) FF(j, n) init[i][j] = ss[i][j];
}

void matmul(int a[][M], int b[][M], int n)
{
    int tp[M][M] = {0};
    FF(i, n) FF(k, n) if(a[i][k]) FF(j, n) if(b[k][j])
        tp[i][j] = (tp[i][j] + a[i][k]*b[k][j]) % mod;
    FF(i, n) FF(j, n) a[i][j] = tp[i][j];
}

void matmul(int a[], int b[][M], int n)
{
    int tp[M] = {0};
    FF(j, n) if(a[j]) FF(i, n) if(b[i][j])
        tp[i] = (tp[i] + b[i][j]*a[j]) % mod;
    FF(i, n) a[i] = tp[i];
}

void qmod(int n, int b)
{
    FF(i, n) FF(j, n) ret[i][j] = (i==j);
    for ( ; b; b >>= 1)
    {
        if (b & 1) matmul(ret, init, n);
        matmul(init, init, n);
    }
}

int main()
{
    int n, x, y;
    while (cin >> n >> x >> y)
    {
        FF(i, M) ans[i] = 1;
        ini(M);
        init[1][1] = (x%mod)*(x%mod) % mod;
        init[1][2] = (y%mod)*(y%mod) % mod;
        init[1][3] = (x%mod)*(y%mod) % mod * 2 % mod;
        init[3][1] = x % mod;
        init[3][3] = y % mod;
        qmod(M, n);
        matmul(ans, ret, M);
        cout << ans[0] << endl;
    }
    return 0;
}