快速幂运算在动态规划中的应用

Posted 2021-04-18 花花酱LeetCode

我们知道动态规划主要可以用来解决两类问题：

1. 优化问题

2. 组合计数问题

花花今天就来和大家聊一聊在组合计数问题中，如何使用快速幂运算来进行加速，把时间复杂度从O(n)降低到O(logn)。用到的知识点有二进制和矩阵运算。

在这期节目中我们假设大家已经推导出了组合计数问题的状态转移方程，如果对动态规划不太熟悉的同学们可以翻一下往期的文章，或者回复 DP  获取内容。

首先我们来看一下  LeetCode 1411. Number of Ways to Paint N × 3 Grid，（回复 1411 获取文章/讲解视频链接）

状态转移方程：

dp[i][0] = dp[i - 1][0] * 3 + dp[i - 1][1] * 2

dp[i][1] = dp[i - 1][0] * 2 + dp[i - 1][1] * 2

边界条件：

dp[1][0] = dp[1][1] = 6

求：

dp[n][0] + dp[n][1]

对于这样的转移方程，我们一般使用递推去求解，一个for循环搞定。时间复杂度O(n)，空间复杂度O(n)，由于dp[i]只和dp[i-1]有关，可以进一步使用滚动数组/临时变量来降维，把空间复杂度降低到O(1)。

很多同学以为线性的递推就是最快的了，其实不然。递推虽然实现简单，也能通过所有测试样例，但并不是最优的。我们还可以做的更好，只需要把状态转移方程改写成矩阵形式即可：

我们可以直接写出dp[n]和dp[1]的关系：

这样问题就变成如何快速求解常数矩阵的n次幂。

我们今天就来介绍一种快速整数幂运算的方法：平方法。

一般形式为 x = a*b^n，底数b可以是实数也可以是矩阵，n必须是正整数。我们从b开始，每次把b平方，就能得到b^(2^k)。在这个过程中，把n看成一个二进制数，如果第k位为1，表示最终结果需要乘以b^(2^k)，否则就跳过。

我们以x = 2 * 3^43为例，43的二进制为：0b101011，第2位和第4位位0，也就是要跳过3^(2^2) = 3^4 和 3^(2^4) = 3^16，我们得到：

快速幂运算的伪代码如下：

Note1：如果a，b为矩阵，则把乘法的部分换成矩阵相乘即可。
Note2：由于结果会非常大，通常每一步计算后都需要进行取余操作。

很明显时间复杂度是O(logn)，因为n的二进制长度为O(logn)，循环会执行O(logn)次，循环内所有的操作都是O(1)，所以总的时间复杂度是O(logn)。

回到我们之前说的1411题，dp[n] = dp[0] * M^(n-1) 

C++代码如下，其中mul是最简单的矩阵相乘的实现。

我们再来看另外一个例子：LeetCode 935. Knight Dialer，这次的系数矩阵很大，10x10，但之后的代码是一样的。

这两题的dp[i]只和dp[i-1]有关，改写成矩阵形式比较简单。我们再来看两个稍微复杂一些的例子。

LeetCode 70. Climbing Stairs 这道题就是让你求斐波那契数列的第n项。

转移方程：dp[i] = dp[i-1] + dp[i-2]

边界条件：dp[0] = dp[1] = 1

同样我们可以把它改写成矩阵形式，需要构造一个2x2的矩阵

可以得到：

其中dp[-1] = 0

后面的计算部分是一样的，最后我们的答案就存储在dp[0][0]。

Note：这题不需要取余。

最后一个例子是 790. Domino and Tromino Tiling

状态转移方程比较复杂：

f[i] = f[i-1] + f[i-2] + 2 * g[i-1]

g[i] = f[i-2] + g[i-1]

边界条件：

f[0] = f[1] = 1, g[0] = g[1] = 0

求：

f[n]

我们依旧可以把它改写成矩阵形式，变量矩阵是1x3，系数矩阵是3x3：

可以得到：

最后答案就存储在dp[0][0]中，参考代码如下：

写了这么多，相信大家对快速幂运算以及其在动态规划中的应用有了一定的了解，不要求掌握，感兴趣的同学可以去自己推导和实现一下。