Largest Submatrix With Rearrangements

lcpy gen -n 1727   # by problem number
lcpy gen -s largest_submatrix_with_rearrangements   # by problem name

Input: matrix = [[0,0,1],[1,1,1],[1,0,1]]
Output: 4
Explanation: You can rearrange the columns as shown above.
The largest submatrix of 1s, in bold, has an area of 4.

Input: matrix = [[1,0,1,0,1]]
Output: 3
Explanation: You can rearrange the columns as shown above.
The largest submatrix of 1s, in bold, has an area of 3.

Input: matrix = [[1,1,0],[1,0,1]]
Output: 2
Explanation: Notice that you must rearrange entire columns, and there is no way to make a submatrix of 1s larger than an area of 2.

class Solution:
    # Time: O(m * n log n)
    # Space: O(n)
    def largest_submatrix(self, matrix: list[list[int]]) -> int:
        n = len(matrix[0])
        heights = [0] * n
        best = 0
        for row in matrix:
            for j, val in enumerate(row):
                heights[j] = heights[j] + 1 if val else 0
            sorted_heights = sorted(heights, reverse=True)
            for i, h in enumerate(sorted_heights):
                best = max(best, h * (i + 1))
        return best