feat: Strassen's matrix multiplication algorithm added

Author: sourav-625

matrix/strassen_matrix_multiply.py

@@ -14,122 +14,173 @@
 
 Matrix = list[list[int]]
 
-def add(A: Matrix, B: Matrix) -> Matrix:
-    n = len(A)
-    return [[A[i][j] + B[i][j] for j in range(n)] for i in range(n)]
 
+def add(matrix_a: Matrix, matrix_b: Matrix) -> Matrix:
+    """
+    Add two square matrices of the same size.
+
+    >>> add([[1,2],[3,4]], [[5,6],[7,8]])
+    [[6, 8], [10, 12]]
+    """
+    n = len(matrix_a)
+    return [[matrix_a[i][j] + matrix_b[i][j] for j in range(n)] for i in range(n)]
+
+
+def sub(matrix_a: Matrix, matrix_b: Matrix) -> Matrix:
+    """
+    Subtract matrix_b from matrix_a.
+
+    >>> sub([[5,6],[7,8]], [[1,2],[3,4]])
+    [[4, 4], [4, 4]]
+    """
+    n = len(matrix_a)
+    return [[matrix_a[i][j] - matrix_b[i][j] for j in range(n)] for i in range(n)]
 
-def sub(A: Matrix, B: Matrix) -> Matrix:
-    n = len(A)
-    return [[A[i][j] - B[i][j] for j in range(n)] for i in range(n)]
 
+def naive_mul(matrix_a: Matrix, matrix_b: Matrix) -> Matrix:
+    """
+    Multiply two square matrices using the naive O(n^3) method.
 
-def naive_mul(A: Matrix, B: Matrix) -> Matrix:
-    n = len(A)
-    C = [[0] * n for _ in range(n)]
+    >>> naive_mul([[1,2],[3,4]], [[5,6],[7,8]])
+    [[19, 22], [43, 50]]
+    """
+    n = len(matrix_a)
+    result = [[0] * n for _ in range(n)]
     for i in range(n):
-        ai = A[i]
-        ci = C[i]
+        row_a = matrix_a[i]
+        row_result = result[i]
         for k in range(n):
-            a_ik = ai[k]
-            bk = B[k]
+            a_ik = row_a[k]
+            col_b = matrix_b[k]
             for j in range(n):
-                ci[j] += a_ik * bk[j]
-    return C
+                row_result[j] += a_ik * col_b[j]
+    return result
 
 
 def next_power_of_two(n: int) -> int:
-    p = 1
-    while p < n:
-        p <<= 1
-    return p
+    """
+    Return the next power of two greater than or equal to n.
+
+    >>> next_power_of_two(5)
+    8
+    """
+    power = 1
+    while power < n:
+        power <<= 1
+    return power
+
 
+def pad_matrix(matrix: Matrix, size: int) -> Matrix:
+    """
+    Pad a matrix with zeros to reach the given size.
 
-def pad_matrix(A: Matrix, size: int) -> Matrix:
-    n = len(A)
+    >>> pad_matrix([[1,2],[3,4]], 4)
+    [[1, 2, 0, 0], [3, 4, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]
+    """
+    rows = len(matrix)
+    cols = len(matrix[0])
     padded = [[0] * size for _ in range(size)]
-    for i in range(n):
-        for j in range(len(A[0])):
-            padded[i][j] = A[i][j]
+    for i in range(rows):
+        for j in range(cols):
+            padded[i][j] = matrix[i][j]
     return padded
 
 
-def unpad_matrix(A: Matrix, rows: int, cols: int) -> Matrix:
-    return [row[:cols] for row in A[:rows]]
+def unpad_matrix(matrix: Matrix, rows: int, cols: int) -> Matrix:
+    """
+    Remove padding from a matrix.
+
+    >>> unpad_matrix([[1,2,0],[3,4,0],[0,0,0]], 2, 2)
+    [[1, 2], [3, 4]]
+    """
+    return [row[:cols] for row in matrix[:rows]]
+
 
+def split(matrix: Matrix) -> tuple:
+    """
+    Split a matrix into four quadrants (top-left, top-right, bottom-left, bottom-right).
 
-def split(A: Matrix) -> tuple:
-    n = len(A)
+    >>> split([[1,2],[3,4]])
+    ([[1]], [[2]], [[3]], [[4]])
+    """
+    n = len(matrix)
     mid = n // 2
-    A11 = [[A[i][j] for j in range(mid)] for i in range(mid)]
-    A12 = [[A[i][j] for j in range(mid, n)] for i in range(mid)]
-    A21 = [[A[i][j] for j in range(mid)] for i in range(mid, n)]
-    A22 = [[A[i][j] for j in range(mid, n)] for i in range(mid, n)]
-    return A11, A12, A21, A22
+    top_left = [[matrix[i][j] for j in range(mid)] for i in range(mid)]
+    top_right = [[matrix[i][j] for j in range(mid, n)] for i in range(mid)]
+    bottom_left = [[matrix[i][j] for j in range(mid)] for i in range(mid, n)]
+    bottom_right = [[matrix[i][j] for j in range(mid, n)] for i in range(mid, n)]
+    return top_left, top_right, bottom_left, bottom_right
 
 
-def join(C11: Matrix, C12: Matrix, C21: Matrix, C22: Matrix) -> Matrix:
-    n2 = len(C11)
+def join(c11: Matrix, c12: Matrix, c21: Matrix, c22: Matrix) -> Matrix:
+    """
+    Join four quadrants into a single matrix.
+
+    >>> join([[1]], [[2]], [[3]], [[4]])
+    [[1, 2], [3, 4]]
+    """
+    n2 = len(c11)
     n = n2 * 2
-    C = [[0] * n for _ in range(n)]
+    result = [[0] * n for _ in range(n)]
     for i in range(n2):
         for j in range(n2):
-            C[i][j] = C11[i][j]
-            C[i][j + n2] = C12[i][j]
-            C[i + n2][j] = C21[i][j]
-            C[i + n2][j + n2] = C22[i][j]
-    return C
+            result[i][j] = c11[i][j]
+            result[i][j + n2] = c12[i][j]
+            result[i + n2][j] = c21[i][j]
+            result[i + n2][j + n2] = c22[i][j]
+    return result
 
 
-def strassen(A: Matrix, B: Matrix, threshold: int = 64) -> Matrix:
+def strassen(matrix_a: Matrix, matrix_b: Matrix, threshold: int = 64) -> Matrix:
     """
-    Multiply square matrices A and B using Strassen algorithm.
-    threshold: below this size, uses naive multiplication (tweakable).
+    Multiply two square matrices using Strassen's algorithm.
+    Uses naive multiplication for matrices smaller than threshold.
+
+    >>> strassen([[1,2],[3,4]], [[5,6],[7,8]])
+    [[19, 22], [43, 50]]
     """
-    assert len(A) == len(A[0]) == len(B) == len(B[0]), (
-        "Only square matrices supported in this implementation"
+    assert len(matrix_a) == len(matrix_a[0]) == len(matrix_b) == len(matrix_b[0]), (
+        "Only square matrices supported"
     )
 
-    n_orig = len(A)
+    n_orig = len(matrix_a)
     if n_orig == 0:
         return []
 
     if (m := next_power_of_two(n_orig)) != n_orig:
-        A_pad = pad_matrix(A, m)
-        B_pad = pad_matrix(B, m)
+        a_pad = pad_matrix(matrix_a, m)
+        b_pad = pad_matrix(matrix_b, m)
     else:
-        A_pad, B_pad = A, B
-
-    C_pad = _strassen_recursive(A_pad, B_pad, threshold)
+        a_pad, b_pad = matrix_a, matrix_b
 
-    C = unpad_matrix(C_pad, n_orig, n_orig)
-    return C
+    c_pad = _strassen_recursive(a_pad, b_pad, threshold)
+    return unpad_matrix(c_pad, n_orig, n_orig)
 
 
-def _strassen_recursive(A: Matrix, B: Matrix, threshold: int) -> Matrix:
-    n = len(A)
+def _strassen_recursive(matrix_a: Matrix, matrix_b: Matrix, threshold: int) -> Matrix:
+    n = len(matrix_a)
     if n <= threshold:
-        return naive_mul(A, B)
+        return naive_mul(matrix_a, matrix_b)
     if n == 1:
-        return [[A[0][0] * B[0][0]]]
+        return [[matrix_a[0][0] * matrix_b[0][0]]]
 
-    A11, A12, A21, A22 = split(A)
-    B11, B12, B21, B22 = split(B)
+    a11, a12, a21, a22 = split(matrix_a)
+    b11, b12, b21, b22 = split(matrix_b)
 
-    M1 = _strassen_recursive(add(A11, A22), add(B11, B22), threshold)
-    M2 = _strassen_recursive(add(A21, A22), B11, threshold)
-    M3 = _strassen_recursive(A11, sub(B12, B22), threshold)
-    M4 = _strassen_recursive(A22, sub(B21, B11), threshold)
-    M5 = _strassen_recursive(add(A11, A12), B22, threshold)
-    M6 = _strassen_recursive(sub(A21, A11), add(B11, B12), threshold)
-    M7 = _strassen_recursive(sub(A12, A22), add(B21, B22), threshold)
+    m1 = _strassen_recursive(add(a11, a22), add(b11, b22), threshold)
+    m2 = _strassen_recursive(add(a21, a22), b11, threshold)
+    m3 = _strassen_recursive(a11, sub(b12, b22), threshold)
+    m4 = _strassen_recursive(a22, sub(b21, b11), threshold)
+    m5 = _strassen_recursive(add(a11, a12), b22, threshold)
+    m6 = _strassen_recursive(sub(a21, a11), add(b11, b12), threshold)
+    m7 = _strassen_recursive(sub(a12, a22), add(b21, b22), threshold)
 
-    C11 = add(sub(add(M1, M4), M5), M7)
-    C12 = add(M3, M5)
-    C21 = add(M2, M4)
-    C22 = add(sub(add(M1, M3), M2), M6)
+    c11 = add(sub(add(m1, m4), m5), m7)
+    c12 = add(m3, m5)
+    c21 = add(m2, m4)
+    c22 = add(sub(add(m1, m3), m2), m6)
 
-    return join(C11, C12, C21, C22)
+    return join(c11, c12, c21, c22)
 
 
 if __name__ == "__main__":
@@ -141,7 +192,6 @@ def _strassen_recursive(A: Matrix, B: Matrix, threshold: int) -> Matrix:
     for row in C:
         print(row)
 
-    # verify against naive
     expected = naive_mul(A, B)
     assert C == expected, "Strassen result differs from naive multiplication!"
     print("Verified: result matches naive multiplication.")