【分治算法】【Python实现】Strassen矩阵乘法

个人主页：丷从心·

系列专栏：分治算法

学习指南：算法学习指南

---

问题描述

• 设$A$和$B$是两个$n\times n$矩阵，$A$和$B$的乘积矩阵$C$中元素$c_{ij}=\sum _{k=1}^n{a}_{ik}{b}_{kj}$
• 每计算$C$的一个元素$c_{ij}$，需要做$n$次乘法和$n-1$次加法，求出矩阵$C$的$n^2$个元素所需的时间为$O(n^3)$

---

基础算法

• 假设$n$是$2$的幂，将矩阵$A$、$B$和$C$中每个矩阵都分块成$4$个大小相等的子矩阵，每个子矩阵都是$n/2\times n/2$的方阵

$$\left|\begin{array}{cc}{C}_{11}& C_{12}\\ C_{21}& C_{22}\end{array}\right|=\left|\begin{array}{cc}{A}_{11}& A_{12}\\ A_{21}& A_{22}\end{array}\right|\left|\begin{array}{cc}{B}_{11}& B_{12}\\ B_{21}& B_{22}\end{array}\right|$$

$$\begin{gathered} C_{11}=A_{11}{B}_{11}+A_{12}{B}_{21} \\ {C}_{12}=A_{11}{B}_{12}+A_{12}{B}_{22} \\ {C}_{21}=A_{21}{B}_{11}+A_{22}{B}_{21} \\ {C}_{22}=A_{21}{B}_{12}+A_{22}{B}_{22} \end{gathered}$$

时间复杂性

• 计算$2$个$n$阶方阵的乘积转化为计算$8$个$n/2$阶方阵的乘积和$4$个$n/2$阶方阵的加法，$2$个$n/2\times n/2$矩阵的加法显然可以在$O(n^2)$时间内完成

$$T(n)=\{\begin{array}{ll}O(1)& n=2\\ 8T(n/2)+O(n^2)& n>2\end{array}$$

$$T(n)=O(n^3)$$

---

Strassen算法

• Strassen算法只用了$7$次乘法运算，但增加了加减法的运算次数

$$\begin{gathered} M_1=A_{11}(B_{12}-B_{22}) \\ {M}_2=(A_{11}+A_{12})B_{22} \\ {M}_3=(A_{21}+A_{22})B_{11} \\ {M}_4=A_{22}(B_{21}-B_{11}) \\ {M}_5=(A_{11}+A_{22})(B_{11}+B_{22}) \\ {M}_6=(A_{12}-A_{22})(B_{21}+B_{22}) \\ {M}_7=(A_{11}-A_{21})(B_{11}+B_{12}) \end{gathered}$$

$$\begin{gathered} C_{11}=M_5+M_4-M_2+M_6 \\ {C}_{12}=M_1+M_2 \\ {C}_{21}=M_3+M_4 \\ {C}_{22}=M_5+M_1-M_3-M_7 \end{gathered}$$

时间复杂性

• Strassen算法用了$7$次对于$n/2$阶矩阵乘积的递归调用和$18$次$n/2$阶矩阵的加减运算

$$T(n)=\{\begin{array}{ll}O(1)& n=2\\ 7T(n/2)+O(n^2)& n>2\end{array}$$

$$T(n)=O(n^{\log 7})\approx O(n^{2.81})$$

---

问题时间复杂性

• $Hopcroft$和$Kerr$已经证明：计算$2$个$2\times 2$矩阵的乘积，$7$次乘法是必要的
• 目前最好的计算时间上界是$O(n^{2.376})$，所知的矩阵乘法的最好下界仍是它的平凡下界$\Omega (n^2)$

---

Python实现

import numpy as np


def strassen_matrix_multiply(a, b):
    n = a.shape[0]

    # 如果输入矩阵的维度小于等于阈值, 使用传统的矩阵乘法
    if n <= 128:
        return np.dot(a, b)

    # 将输入矩阵划分为四个子矩阵
    mid = n // 2

    a11 = a[:mid, :mid]
    a12 = a[:mid, mid:]
    a21 = a[mid:, :mid]
    a22 = a[mid:, mid:]

    b11 = b[:mid, :mid]
    b12 = b[:mid, mid:]
    b21 = b[mid:, :mid]
    b22 = b[mid:, mid:]

    # 递归计算七个矩阵乘法
    m1 = strassen_matrix_multiply(a11, b12 - b22)
    m2 = strassen_matrix_multiply(a11 + a12, b22)
    m3 = strassen_matrix_multiply(a21 + a22, b11)
    m4 = strassen_matrix_multiply(a22, b21 - b11)
    m5 = strassen_matrix_multiply(a11 + a22, b11 + b22)
    m6 = strassen_matrix_multiply(a12 - a22, b21 + b22)
    m7 = strassen_matrix_multiply(a11 - a21, b11 + b12)

    # 计算结果矩阵的四个子矩阵
    c11 = m5 + m4 - m2 + m6
    c12 = m1 + m2
    c21 = m3 + m4
    c22 = m5 + m1 - m3 - m7

    # 组合四个子矩阵形成结果矩阵
    c = np.zeros((n, n))

    c[:mid, :mid] = c11
    c[:mid, mid:] = c12
    c[mid:, :mid] = c21
    c[mid:, mid:] = c22

    return c


a = np.random.randint(0, 10, (256, 256))
b = np.random.randint(0, 10, (256, 256))

res = strassen_matrix_multiply(a, b)

print('矩阵 a:')
print(a)

print('\n矩阵 b:')
print(b)

print('\n乘积矩阵 c:')
print(res)

矩阵 a:
[[4 8 5 ... 6 3 7]
 [8 0 0 ... 8 4 6]
 [5 4 2 ... 7 4 8]
 ...
 [6 0 7 ... 5 5 6]
 [9 3 6 ... 6 2 9]
 [2 2 5 ... 7 3 6]]

矩阵 b:
[[8 6 5 ... 3 7 2]
 [2 9 5 ... 4 2 3]
 [1 0 0 ... 0 5 4]
 ...
 [0 6 4 ... 3 9 1]
 [9 9 7 ... 5 4 1]
 [6 0 1 ... 2 6 1]]

乘积矩阵 c:
[[5306. 5218. 5339. ... 5510. 5653. 5120.]
 [5327. 5038. 5228. ... 5132. 5251. 4969.]
 [5325. 5265. 5237. ... 5252. 5474. 4962.]
 ...
 [5491. 5176. 5190. ... 5443. 5463. 5040.]
 [5648. 5509. 5191. ... 5439. 5707. 5273.]
 [5034. 5027. 5061. ... 4866. 5366. 4772.]]

---