diff --git a/DIRECTORY.md b/DIRECTORY.md index 7dcb1f9c2387..67e701ffcf3b 100644 --- a/DIRECTORY.md +++ b/DIRECTORY.md @@ -277,6 +277,7 @@ * [Prefix Sum](data_structures/arrays/prefix_sum.py) * [Product Sum](data_structures/arrays/product_sum.py) * [Rotate Array](data_structures/arrays/rotate_array.py) + * [Set Matrix Zeroes](data_structures/arrays/set_matrix_zeroes.py) * [Sparse Table](data_structures/arrays/sparse_table.py) * [Sudoku Solver](data_structures/arrays/sudoku_solver.py) * Binary Tree diff --git a/data_structures/arrays/set_matrix_zeroes.py b/data_structures/arrays/set_matrix_zeroes.py new file mode 100644 index 000000000000..905d03b2daf5 --- /dev/null +++ b/data_structures/arrays/set_matrix_zeroes.py @@ -0,0 +1,67 @@ +""" +Set Matrix Zeroes Algorithm +--------------------------- +If an element in an m x n matrix is 0, set its entire row and column to 0. + +Explanation: +We use the first row and first column as markers to track which rows and +columns should be zeroed, avoiding extra space usage (O(1) space complexity). + +References: +https://leetcode.com/problems/set-matrix-zeroes/ + +Doctest: +>>> matrix = [ +... [1, 1, 1], +... [1, 0, 1], +... [1, 1, 1] +... ] +>>> set_matrix_zeroes(matrix) +>>> matrix +[[1, 0, 1], [0, 0, 0], [1, 0, 1]] +>>> matrix = [[0, 1, 2, 0], [3, 4, 5, 2], [1, 3, 1, 5]] +>>> set_matrix_zeroes(matrix) +>>> matrix +[[0, 0, 0, 0], [0, 4, 5, 0], [0, 3, 1, 0]] +""" + + +def set_matrix_zeroes(matrix: list[list[int]]) -> None: + """ + Modify the matrix in-place such that if an element is 0, + its entire row and column are set to 0. + + :param matrix: 2D list of integers + :return: None (modifies matrix in-place) + + Time Complexity: O(m * n) + Space Complexity: O(1) + """ + rows = len(matrix) + cols = len(matrix[0]) + col0 = 1 + + # Step 1: Mark rows and columns that need to be zeroed + for i in range(rows): + if matrix[i][0] == 0: + col0 = 0 + for j in range(1, cols): + if matrix[i][j] == 0: + matrix[i][0] = 0 + matrix[0][j] = 0 + + # Step 2: Update the inner matrix cells + for i in range(1, rows): + for j in range(1, cols): + if matrix[i][0] == 0 or matrix[0][j] == 0: + matrix[i][j] = 0 + + # Step 3: Handle the first row + if matrix[0][0] == 0: + for j in range(cols): + matrix[0][j] = 0 + + # Step 4: Handle the first column + if col0 == 0: + for i in range(rows): + matrix[i][0] = 0