From f40df023bc3a5de82c8aaf8ef13c8cda2194d102 Mon Sep 17 00:00:00 2001 From: Nihal Patel Date: Fri, 17 Oct 2025 23:51:24 +0530 Subject: [PATCH] Implement Strassen's matrix multiplication algorithm This file implements Strassen's matrix multiplication algorithm, which is faster than the standard O(n^3) method for large matrices. It includes helper functions for matrix operations and benchmarks the performance of Strassen's algorithm against standard multiplication. --- matrix/strassen.py | 252 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 252 insertions(+) create mode 100644 matrix/strassen.py diff --git a/matrix/strassen.py b/matrix/strassen.py new file mode 100644 index 000000000000..e7f4f7759ead --- /dev/null +++ b/matrix/strassen.py @@ -0,0 +1,252 @@ +""" +Implementation of Strassen's matrix multiplication algorithm. +https://en.wikipedia.org/wiki/Strassen_algorithm + +This is a divide-and-conquer algorithm that is asymptotically faster +than the standard O(n^3) matrix multiplication for large matrices. + +Note: In Python, due to the overhead of recursion and list slicing, +this implementation will be *slower* than the iterative version +for small or medium-sized matrices (like 4x4). +""" + +# type Matrix = list[list[int]] # psf/black currently fails on this line +Matrix = list[list[int]] + +# --- Test Matrices (reused from other files) --- +matrix_1_to_4 = [ + [1, 2], + [3, 4], +] + +matrix_5_to_8 = [ + [5, 6], + [7, 8], +] + +matrix_count_up = [ + [1, 2, 3, 4], + [5, 6, 7, 8], + [9, 10, 11, 12], + [13, 14, 15, 16], +] + +matrix_unordered = [ + [5, 8, 1, 2], + [6, 7, 3, 0], + [4, 5, 9, 1], + [2, 6, 10, 14], +] + +matrix_non_square = [ + [1, 2, 3], + [4, 5, 6], +] + + +# --- Helper function from matrix_multiplication_recursion.py --- +def is_square(matrix: Matrix) -> bool: + """ + Checks if a matrix is square. + >>> is_square(matrix_1_to_4) + True + >>> is_square(matrix_non_square) + False + """ + len_matrix = len(matrix) + return all(len(row) == len_matrix for row in matrix) + + +# --- Helper function for benchmarking --- +def matrix_multiply(matrix_a: Matrix, matrix_b: Matrix) -> Matrix: + """ + Standard iterative matrix multiplication for comparison. + >>> matrix_multiply(matrix_1_to_4, matrix_5_to_8) + [[19, 22], [43, 50]] + """ + return [ + [sum(a * b for a, b in zip(row, col)) for col in zip(*matrix_b)] + for row in matrix_a + ] + + +# --- Helper functions for Strassen's Algorithm --- + + +def matrix_add(matrix_a: Matrix, matrix_b: Matrix) -> Matrix: + """ + Adds two matrices element-wise. + >>> matrix_add(matrix_1_to_4, matrix_5_to_8) + [[6, 8], [10, 12]] + """ + return [ + [matrix_a[i][j] + matrix_b[i][j] for j in range(len(matrix_a[0]))] + for i in range(len(matrix_a)) + ] + + +def matrix_subtract(matrix_a: Matrix, matrix_b: Matrix) -> Matrix: + """ + Subtracts matrix_b from matrix_a element-wise. + >>> matrix_subtract(matrix_5_to_8, matrix_1_to_4) + [[4, 4], [4, 4]] + """ + return [ + [matrix_a[i][j] - matrix_b[i][j] for j in range(len(matrix_a[0]))] + for i in range(len(matrix_a)) + ] + + +def split_matrix(matrix: Matrix) -> tuple[Matrix, Matrix, Matrix, Matrix]: + """ + Splits a given matrix into four equal quadrants. + >>> a, b, c, d = split_matrix(matrix_count_up) + >>> a + [[1, 2], [5, 6]] + >>> b + [[3, 4], [7, 8]] + >>> c + [[9, 10], [13, 14]] + >>> d + [[11, 12], [15, 16]] + """ + n = len(matrix) // 2 + a11 = [row[:n] for row in matrix[:n]] + a12 = [row[n:] for row in matrix[:n]] + a21 = [row[:n] for row in matrix[n:]] + a22 = [row[n:] for row in matrix[n:]] + return a11, a12, a21, a22 + + +def combine_matrices( + c11: Matrix, c12: Matrix, c21: Matrix, c22: Matrix +) -> Matrix: + """ + Combines four quadrants into a single matrix. + >>> a, b, c, d = split_matrix(matrix_count_up) + >>> combine_matrices(a, b, c, d) == matrix_count_up + True + """ + n = len(c11) + result = [] + for i in range(n): + result.append(c11[i] + c12[i]) + for i in range(n): + result.append(c21[i] + c22[i]) + return result + + +def pad_matrix(matrix: Matrix, target_size: int) -> Matrix: + """Pads a matrix with zeros to reach the target_size.""" + n = len(matrix) + if n == target_size: + return matrix + + padded_matrix = [[0] * target_size for _ in range(target_size)] + for i in range(n): + for j in range(len(matrix[i])): + padded_matrix[i][j] = matrix[i][j] + return padded_matrix + + +def unpad_matrix(matrix: Matrix, original_size: int) -> Matrix: + """Removes padding to return to the original_size.""" + if len(matrix) == original_size: + return matrix + return [row[:original_size] for row in matrix[:original_size]] + + +# --- Main Strassen Function --- + + +def strassen(matrix_a: Matrix, matrix_b: Matrix) -> Matrix: + """ + :param matrix_a: A square Matrix. + :param matrix_b: Another square Matrix with the same dimensions as matrix_a. + :return: Result of matrix_a * matrix_b. + :raises ValueError: If the matrices cannot be multiplied. + + >>> strassen([], []) + [] + >>> strassen(matrix_1_to_4, matrix_5_to_8) + [[19, 22], [43, 50]] + >>> strassen(matrix_count_up, matrix_unordered) + [[37, 61, 74, 61], [105, 165, 166, 129], [173, 269, 258, 197], [241, 373, 350, 265]] + >>> strassen(matrix_1_to_4, matrix_non_square) + Traceback (most recent call last): + ... + ValueError: Matrices must be square and of the same dimensions + >>> strassen(matrix_1_to_4, matrix_count_up) + Traceback (most recent call last): + ... + ValueError: Matrices must be square and of the same dimensions + """ + if not matrix_a or not matrix_b: + return [] + + if not ( + len(matrix_a) == len(matrix_b) + and is_square(matrix_a) + and is_square(matrix_b) + ): + raise ValueError("Matrices must be square and of the same dimensions") + + original_size = len(matrix_a) + + # Base case + if original_size == 1: + return [[matrix_a[0][0] * matrix_b[0][0]]] + + # Pad matrix to the next power of 2 + n = original_size + if n & (n - 1) != 0: + next_power_of_2 = 1 << n.bit_length() + a = pad_matrix(matrix_a, next_power_of_2) + b = pad_matrix(matrix_b, next_power_of_2) + n = next_power_of_2 + else: + a = matrix_a + b = matrix_b + + # Split matrices into quadrants + a11, a12, a21, a22 = split_matrix(a) + b11, b12, b21, b22 = split_matrix(b) + + # Calculate the 7 Strassen products recursively + p1 = strassen(a11, matrix_subtract(b12, b22)) + p2 = strassen(matrix_add(a11, a12), b22) + p3 = strassen(matrix_add(a21, a22), b11) + p4 = strassen(a22, matrix_subtract(b21, b11)) + p5 = strassen(matrix_add(a11, a22), matrix_add(b11, b22)) + p6 = strassen(matrix_subtract(a12, a22), matrix_add(b21, b22)) + p7 = strassen(matrix_subtract(a11, a21), matrix_add(b11, b12)) + + # Calculate result quadrants + c11 = matrix_add(matrix_subtract(matrix_add(p5, p4), p2), p6) + c12 = matrix_add(p1, p2) + c21 = matrix_add(p3, p4) + c22 = matrix_subtract(matrix_subtract(matrix_add(p5, p1), p3), p7) + + # Combine result quadrants + result = combine_matrices(c11, c12, c21, c22) + + # Unpad the result to match original dimensions + return unpad_matrix(result, original_size) + + +if __name__ == "__main__": + from doctest import testmod + + failure_count, test_count = testmod() + if not failure_count: + print("\nBenchmark (Note: Strassen has high overhead in Python):") + from functools import partial + from timeit import timeit + + # Run fewer iterations as Strassen is slower for small matrices in Python + mytimeit = partial(timeit, globals=globals(), number=10_0DENIED) + for func in ("matrix_multiply", "strassen"): + print( + f"{func:>25}(): " + f"{mytimeit(f'{func}(matrix_count_up, matrix_unordered)')}" + )