求第n项的因子数量

最近笔试期间遇到一个难题,现在终于解决了,感谢各路大佬的指点,我在这里分享一下结果。

小红拿到一个数列满足:

f(1) = a;      f(2) = b;      f(i) = f(i-1) * f(i-2) * c^d

题目要求计算出第n项的因子数量,因子数量对 10^9+1 取模。

输入:a, b, c, d, n 5个整数, (1 <= a, b, c, d, n <= 10^12)

例如:输入:1 2 3 4 3

           输出:10

解题理论准备:快速幂矩阵、因子与质因子关系

1、因子与质因子的关系-CSDN博客

 2、快速幂算法-python-CSDN博客

3、算法学习笔记(4):快速幂 - 知乎 

python代码如下:

'''
解题思路:
1、分别算出 a, b, c三个数的质因数;
2、通过快速幂矩阵计算出第n项数据中a、b、c的指数(计算过程中要取模);
3、结合a、b、c的指数,以及a, b, c三个数的质因数,
   来求出第n项数据对应的质因数,以及对应质因数的指数;
4、最后将第n项数据的各个质因数的指数分别+1之后相乘,就得到第n项数据的因子个数。
'''

import time
import numpy as np


def add2dict(dic, n):
    if n in dic:
        dic[n] += 1
    else:
        dic[n] = 1
    return dic


def primeFactors(num):
    factors = {}
    i = 2
    while i * i <= num:
        if num % i == 0:
            num //= i
            factors = add2dict(factors, i)
        else:
            i += 1
    if num > 1:
        factors = add2dict(factors, num)
    return factors


# 快速幂矩阵
def matrixFastPower(matrix, power, num_mod):
    matrix = np.array(matrix) % num_mod
    res = np.eye(matrix.shape[0])
    while power:
        if power & 1:
            res = np.dot(res, matrix) % num_mod
        power >>= 1
        matrix = np.dot(matrix, matrix) % num_mod
    return res


if __name__ == '__main__':

    start_time = time.time()
    a, b, c, d, n = int(1e12), int(1e12), int(1e12), int(1e12), int(1e12)
    # a, b, c, d, n = 1, 2, 3, 4, 6
    num_mod = int(10e9 + 7)
    if n == 1:
        factors_a = primeFactors(a, num_mod)
        sum = 1
        for i in factors_a:
            sum *= (factors_a[i] + 1) % num_mod
    elif n == 2:
        factors_b = primeFactors(b)
        sum = 1
        for i in factors_b:
            sum *= (factors_b[i] + 1) % num_mod
    else:
        factors_a = primeFactors(a)
        factors_b = primeFactors(b)
        factors_c = primeFactors(c)

        A = [[0, 1],
             [1, 1]]
        C = [[0, 1, 0],
             [1, 1, 1],
             [0, 0, 1]]

        A_n = matrixFastPower(A, n - 1, num_mod)
        C_n = matrixFastPower(C, n - 1, num_mod)
        # a_pow = A^(n-1) * [a(1), a(2)]
        a_pow = int(np.dot(A_n, np.array([1, 0]).T)[0]) % num_mod
        # b_pow = A^(n-1) * [b(1), b(2)]
        b_pow = int(np.dot(A_n, np.array([0, 1]).T)[0]) % num_mod
        # c_pow = C^(n-1) * [c(1), c(2), d]
        c_pow = int(np.dot(C_n, np.array([0, 0, d]).T)[0]) % num_mod
        # print(a_pow, b_pow, c_pow)

        for i in factors_a:
            factors_a[i] = ((factors_a[i] % num_mod) * a_pow) % num_mod
        for j in factors_b:
            factors_b[j] = ((factors_b[j] % num_mod) * b_pow) % num_mod
            if j not in factors_a:
                factors_a[j] = factors_b[j] % num_mod
            else:
                factors_a[j] = (factors_a[j] + factors_b[j]) % num_mod
        del factors_b
        for k in factors_c:
            factors_c[k] = ((factors_c[k] % num_mod) * c_pow) % num_mod
            if k not in factors_a:
                factors_a[k] = factors_c[k] % num_mod
            else:
                factors_a[k] = (factors_a[k] + factors_c[k]) % num_mod
        del factors_c
        sum = 1
        for ix in factors_a:
            sum *= (factors_a[ix] + 1) % num_mod

    sum %= num_mod
    print(sum)
    print("Time:", round((time.time() - start_time) * 1000, 2), 'ms')

你可能感兴趣的:(牛客,数据结构与算法,Python精修,算法,数据结构,因子,质因子)