lintcode 948. 序列计数

https://www.lintcode.com/problem/sequence-count/description

求有多少个长为 n 的序列满足如下要求:

• 序列中的数都是不超过 m 的正整数
• 这 n 个数的和是 p 的倍数
• 这 n 个数中至少有一个是质数

返回满足要求的序列个数对 10000007 取模的结果

 

• 1 <= n <= 10^9
• 1 <= m <= 2 x 10^7
• 1 <= p <= 100

您在真实的面试中是否遇到过这个题？  是

题目纠错

样例

样例 1:

输入: n = 2, m = 5, p = 4
输出: 3
解释: 这三个序列是 [2, 2], [1, 3] 和 [3, 1]

样例 2:

输入: n = 2, m = 3, p = 4
输出: 5
解释: 这 5 个序列是 [2, 2], [1, 3], [3, 1], [3, 5], [5, 3]

思路：先求<m的所有合数。至少有一个质数数转化为：所有数的情况 - 所有合数的情况。dp1[i][a]代表前i位时（i<n）可以取所有数时的情况数。count1[b]为<=m时,%p==b的数的个数。对应dp2[a]，count2[b]为合数情况。dp方程为dp1[i+1][s]=∑dp1[i][a]*count1[b]。（所有a+b==s的情况）。因为n过大。需要矩阵快速幂加速。复杂度O（log n * q^3）。矩阵相当于每列由count组成的，其中每一列与前一列 比较 相当于转动了一格。

public class Solution {
    /**
     * @param n: an integer
     * @param m: an integer
     * @param p: an integer
     * @return: return an integer
     */
    long mod=10000007;
    public int sequenceCount(int n, int m, int p) {
        boolean[] rec=new boolean[m+1];
        rec[1]=true;
        for(int i=2;i*i<=m;i++){
            if(rec[i])continue;//true 是合数
            for(int j=i;i*j<=m;j++){
                rec[i*j]=true;
            }
        }
        long[] count1=new long[p];
        long[] count2=new long[p]; 
        for(int i=1;i<=m;i++){
            int r=i%p;
            count1[r]++;
            if(rec[i]){
                count2[r]++;
            }
        }
        long[] dp1=new long[p];
        long[] dp2=new long[p];
        dp1[0]=1;
        dp2[0]=1;
        long[] ret1=work(dp1,count1,n);
        long[] ret2=work(dp2,count2,n);
        long ret=ret1[0]-ret2[0];
        if(ret<0)ret+=mod;
        return (int)ret;
    }
    long[] work(long[] dp,long[] count,long k){
        int n=dp.length;
        long[][] A=new long[n][n];
        for(int j=n-1,cur=0;j>=0;j--,cur++){
            for(int i=0;i<n;i++){
                int index=((n-1-i-cur)+n)%n;
                A[i][j]=count[index];
            }
        }
        while(k>1){
            if(k%2==1){
                dp=multiply1(dp,A);
            }
            A=multiply2(A,A);
            k/=2;
        }
        return multiply1(dp,A);
    }
    long[] multiply1(long[] dp,long[][] A){
        int n=dp.length;
        long[] ret=new long[n];
        for(int i=0;i<n;i++){
            for(int j=0;j<n;j++){
                ret[i]+=dp[j]*A[j][i];
                ret[i]%=mod;
            }
        }
        return ret;
    }
    
    long[][] multiply2(long[][] A,long[][] B){
        int n=A.length;
        int m=B[0].length;
        long[][] ret=new long[n][m];
        for(int i=0;i<n;i++){
            for(int j=0;j<m;j++){
                for(int k=0;k<m;k++){
                    ret[i][j]+=A[i][k]*B[k][j];
                    ret[i][j]%=mod;
                }
            }
        }
        return ret;
    }
}