快速幂算法

今天在Leetcode周赛里碰到这样一个题目：

5802. 统计好数字的数目

我们称一个数字字符串是 好数字 当它满足（下标从 0 开始）偶数 下标处的数字为 偶数 且 奇数 下标处的数字为 质数 （2，3，5 或 7）。

• 比方说，"2582" 是好数字，因为偶数下标处的数字（2 和 8）是偶数且奇数下标处的数字（5 和 2）为质数。但 "3245" 不是 好数字，因为 3 在偶数下标处但不是偶数。

给你一个整数 n ，请你返回长度为 n 且为好数字的数字字符串 总数 。由于答案可能会很大，请你将它对 109+7 取余后返回 。

一个 数字字符串 是每一位都由 1 到 9 组成的字符串，且可能包含前导 0 。

提示：

• 1 <= n <= 1015

根据题目描述，容易发现，偶数位有0、2、4、6、8五个选择，奇数位有2、3、5、7四个选择。

一个长度为n的数字字符串，有$\lfloor \frac{n+1}{2}\rfloor$个偶数位，有$\lfloor \frac{n}{2}\rfloor$个奇数位。

由此可得长度为 n 且为好数字的数字字符串个数计算公式为$5^{\lfloor \frac{n+1}{2}\rfloor }\times 4^{\lfloor \frac{n}{2}\rfloor }$。

思路很简单，但是重点在提示部分，1 <= n <= 1015。如果采用累乘来计算an，时间复杂度为O(n)，在这题显然是行不通的。

快速幂算法

在比赛结束后，查看牛人的题解时，发现他们都是使用自定义的函数来计算幂。搜索后发现这个快速幂算法可以将an的时间复杂度降至为O(log n)。

算法思想摘自快速幂百度百科词条：

快速幂算法的核心思想就是每一步都把指数分成两半，而相应的底数做平方运算。这样不仅能把非常大的指数给不断变小，所需要执行的循环次数也变小，而最后表示的结果却一直不会变。

为了更好地理解这个算法，推导了一遍这个公式：

$需要计算a^n，为了方便理解，假设n=13，则n的二进制表示为1101。$

$n=1\times 2^3+1\times 2^2+0\times 2^1+1\times 2^0$

$a^n=a^{13}=a^{1\times 2^3+1\times 2^2+0\times 2^1+1\times 2^0}=a^{1\times 2^3}\times a^{1\times 2^2}\times a^{0\times 2^1}\times a^{1\times 2^0}$

$其中重点是计算a^{2^i}，又可得a^{2^{i+1}}=a^{2^i\times 2}=a^{2^i+2^i}=a^{2^i}\times a^{2^i}，每次向前计算时乘以自己本身即可。$

由上述公式可写出如下Python代码：

def fast_pow(a, n, mod):
    result = 1
    while n:
        if n & 1:
            result = result * a % mod
        a = a * a % mod
        n >>= 1
    return result

在代码中使用了位运算，在第4行n & 1使用的是与运算，只有两个位都为1时，结果才为1。由于1只有最后一位为1，所以可以获得n最后一位的数值情况。在第7行n >>= 1使用的是移位，即每次把n的二进制数向右移动一位，这样会将最后一位删除，在首位补0。

时间复杂度分析

在代码中，只有一个while循环，所以只需要计算循环的次数即可。这个循环的条件是n，相当于n != 0，改变n的值的是第7行n >>= 1，向右移动一位。当n中所有1都被删除后，跳出循环。所以这个循环运行次数为n二进制的位数，即log n，时间复杂度为O(log n)。

对比

import time


def time_clock(fun):
    def t(a, b, mod):
        s_time = time.time()
        r = fun(a, b, mod)
        print(f'r: {r}')
        e_time = time.time()
        print(f'time: {e_time-s_time:.4f}')

    return t


@time_clock
def fast_pow(a, b, mod):
    result = 1
    while b:
        if b & 1:
            result = result * a % mod
        a = a * a % mod
        b >>= 1
    return result


@time_clock
def my_pow(a, b, mod):
    result = 1
    for i in range(b):
        result = result * a % mod
    return result


a = 2
b = 1000000000
mod = 1000000007
x = my_pow(a, b, mod)
y = fast_pow(a, b, mod)

# 累乘输出
# r: 140625001
# time: 196.3966
# 快速幂输出
# r: 140625001
# time: 0.0000

解题

def fast_pow(a, b, mod):
    result = 1
    while b:
        if b & 1:
            result = result * a % mod
        a = a * a % mod
        b >>= 1
    return result


class Solution:
    def countGoodNumbers(self, n: int) -> int:
        mod = 10**9+7
        a = int((n+1)/2)
        b = int(n/2)
        result = fast_pow(5, a, mod) * fast_pow(4, b, mod) % mod
        return int(result)

本文首发于微信公众号<卡卡KSCN>，欢迎关注