Merge branch 'strassen_matrix_multiply' of https://github.com/sourav-625/Python into strassen_matrix_multiply

Author: sourav-625

matrix/strassen_matrix_multiply.py

@@ -12,20 +12,21 @@
 https://en.wikipedia.org/wiki/Strassen_algorithm
 """
 
-
 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 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(A: Matrix, B: Matrix) -> Matrix:
     n = len(A)
-    C = [[0]*n for _ in range(n)]
+    C = [[0] * n for _ in range(n)]
     for i in range(n):
         ai = A[i]
         ci = C[i]
@@ -36,23 +37,27 @@ def naive_mul(A: Matrix, B: Matrix) -> Matrix:
                 ci[j] += a_ik * bk[j]
     return C
 
+
 def next_power_of_two(n: int) -> int:
     p = 1
     while p < n:
         p <<= 1
     return p
 
+
 def pad_matrix(A: Matrix, size: int) -> Matrix:
     n = len(A)
-    padded = [[0]*size for _ in range(size)]
+    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]
     return padded
 
+
 def unpad_matrix(A: Matrix, rows: int, cols: int) -> Matrix:
     return [row[:cols] for row in A[:rows]]
 
+
 def split(A: Matrix) -> tuple:
     n = len(A)
     mid = n // 2
@@ -62,10 +67,11 @@ def split(A: Matrix) -> tuple:
     A22 = [[A[i][j] for j in range(mid, n)] for i in range(mid, n)]
     return A11, A12, A21, A22
 
+
 def join(C11: Matrix, C12: Matrix, C21: Matrix, C22: Matrix) -> Matrix:
     n2 = len(C11)
     n = n2 * 2
-    C = [[0]*n for _ in range(n)]
+    C = [[0] * n for _ in range(n)]
     for i in range(n2):
         for j in range(n2):
             C[i][j] = C11[i][j]
@@ -74,19 +80,21 @@ def join(C11: Matrix, C12: Matrix, C21: Matrix, C22: Matrix) -> Matrix:
             C[i + n2][j + n2] = C22[i][j]
     return C
 
+
 def strassen(A: Matrix, B: Matrix, threshold: int = 64) -> Matrix:
     """
     Multiply square matrices A and B using Strassen algorithm.
     threshold: below this size, uses naive multiplication (tweakable).
     """
-    assert len(A) == len(A[0]) == len(B) == len(B[0]), "Only square matrices supported in this implementation"
+    assert len(A) == len(A[0]) == len(B) == len(B[0]), (
+        "Only square matrices supported in this implementation"
+    )
 
     n_orig = len(A)
     if n_orig == 0:
         return []
 
-    m = next_power_of_two(n_orig)
-    if m != n_orig:
+    if (m := next_power_of_two(n_orig)) != n_orig:
         A_pad = pad_matrix(A, m)
         B_pad = pad_matrix(B, m)
     else:
@@ -97,6 +105,7 @@ def strassen(A: Matrix, B: Matrix, threshold: int = 64) -> Matrix:
     C = unpad_matrix(C_pad, n_orig, n_orig)
     return C
 
+
 def _strassen_recursive(A: Matrix, B: Matrix, threshold: int) -> Matrix:
     n = len(A)
     if n <= threshold:
@@ -122,17 +131,10 @@ def _strassen_recursive(A: Matrix, B: Matrix, threshold: int) -> Matrix:
 
     return join(C11, C12, C21, C22)
 
+
 if __name__ == "__main__":
-    A = [
-        [1, 2, 3],
-        [4, 5, 6],
-        [7, 8, 9]
-    ]
-    B = [
-        [9, 8, 7],
-        [6, 5, 4],
-        [3, 2, 1]
-    ]
+    A = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
+    B = [[9, 8, 7], [6, 5, 4], [3, 2, 1]]
 
     C = strassen(A, B, threshold=1)
     print("A * B =")
@@ -142,4 +144,4 @@ def _strassen_recursive(A: Matrix, B: Matrix, threshold: int) -> Matrix:
     # verify against naive
     expected = naive_mul(A, B)
     assert C == expected, "Strassen result differs from naive multiplication!"
-    print("Verified: result matches naive multiplication.")
+    print("Verified: result matches naive multiplication.")