파이썬으로 이진 행렬에서 1로만 이루어진 정사각형 부분 행렬의 개수 구하기

문제 소개

2차원 이진 행렬(0과 1로만 구성된 행렬)이 주어졌을 때, 모든 원소가 1로 이루어진 정사각형 부분 행렬의 총 개수를 구하는 프로그램을 만들어 보겠습니다.

예를 들어 입력이 다음과 같다고 가정해 봅시다.

이 경우 출력은 17이 됩니다. 크기가 1×1인 정사각형이 12개, 2×2인 정사각형이 4개, 3×3인 정사각형이 1개 존재하기 때문입니다.

풀이 접근 방식: 동적 계획법

이 문제는 동적 계획법(Dynamic Programming)을 활용하면 효율적으로 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다.

특정 위치 (i, j)를 오른쪽 아래 꼭짓점으로 하는 정사각형의 개수는, 그 위치에서 만들 수 있는 가장 큰 정사각형의 한 변의 길이와 같습니다.

예를 들어 어떤 칸의 값이 3으로 갱신되었다면, 그 칸을 꼭짓점으로 하는 1×1, 2×2, 3×3 정사각형이 각각 하나씩 존재한다는 의미이므로 총 3개를 결과에 더할 수 있습니다.

구체적인 알고리즘 단계는 다음과 같습니다.

• 결괏값을 저장할 변수 res를 0으로 초기화합니다.
• 행렬의 각 행 인덱스 i에 대해 반복합니다.
  • 각 열 인덱스 j에 대해 반복합니다.
    • i가 0이거나 j가 0인 경우(첫 번째 행 또는 첫 번째 열): 경계이므로 별도의 계산 없이 res에 matrix[i][j] 값을 그대로 더합니다.
    • 그 외의 경우 matrix[i][j]가 1이라면, 왼쪽(matrix[i][j-1]), 위쪽(matrix[i-1][j]), 왼쪽 위 대각선(matrix[i-1][j-1]) 값 중 최솟값에 1을 더한 값으로 현재 칸을 갱신하고, 갱신된 값을 res에 더합니다.
• 모든 순회가 끝나면 res를 반환합니다.

세 방향 값 중 최솟값에 1을 더하는 이유는, 왼쪽·위쪽·대각선 방향으로 확장 가능한 정사각형의 크기 중 가장 작은 값에 의해 현재 위치에서 만들 수 있는 최대 정사각형 크기가 결정되기 때문입니다.

파이썬 구현 예시

class Solution:
    def solve(self, matrix):
        res = 0
        for i in range(len(matrix)):
            for j in range(len(matrix[0])):
                # 첫 번째 행 또는 첫 번째 열인 경우
                if i == 0 or j == 0:
                    res += matrix[i][j]
                elif matrix[i][j] == 1:
                    # 왼쪽, 위쪽, 왼쪽 위 대각선 중 최솟값 + 1로 갱신
                    matrix[i][j] = min(matrix[i][j - 1], matrix[i - 1][j], matrix[i - 1][j - 1]) + 1
                    res += matrix[i][j]
        return res

ob = Solution()
matrix = [
    [1, 1, 1, 0],
    [1, 1, 1, 0],
    [1, 1, 1, 0],
    [0, 0, 0, 0],
    [1, 0, 1, 1]
]
print(ob.solve(matrix))

입력

matrix = [
[1, 1, 1, 0],
[1, 1, 1, 0],
[1, 1, 1, 0],
[0, 0, 0, 0],
[1, 0, 1, 1]
]

출력

17

복잡도 분석

행렬의 크기를 m×n이라 할 때, 모든 칸을 한 번씩만 방문하므로 시간 복잡도는 O(m×n)입니다. 또한 입력 행렬 자체를 DP 테이블처럼 직접 갱신하여 사용하기 때문에 추가 메모리는 O(1)로 매우 효율적입니다. 원본 행렬을 유지해야 하는 경우에는 별도의 DP 배열을 복제하여 사용하면 됩니다.