파이썬으로 모든 요소가 1인 정사각형 부분 행렬의 개수 구하기

m × n 크기의 이진 행렬(0과 1로만 이루어진 행렬)이 주어졌을 때, 모든 요소가 1로 채워져 있는 정사각형 부분 행렬이 몇 개나 있는지 구하는 문제입니다.

문제 예시

다음과 같은 행렬이 입력으로 주어졌다고 가정해 보겠습니다.

이 경우 출력은 15가 됩니다. 그 이유는 다음과 같습니다.

• 한 변의 길이가 1인 정사각형: 10개
• 한 변의 길이가 2인 정사각형: 4개
• 한 변의 길이가 3인 정사각형: 1개

따라서 전체 정사각형의 개수는 10 + 4 + 1 = 15개입니다.

해결 접근 방법 (동적 계획법)

이 문제는 동적 계획법(Dynamic Programming)으로 효율적으로 해결할 수 있습니다. 핵심 아이디어는 각 위치에서 끝나는 가장 큰 정사각형의 한 변의 길이를 행렬 자체에 저장하는 것입니다.

특정 칸 (row, col)의 값이 1일 때, 그 칸을 오른쪽 아래 꼭짓점으로 하는 정사각형의 최대 크기는 왼쪽, 위쪽, 왼쪽 위 대각선 방향 세 칸의 값 중 최솟값에 1을 더한 값이 됩니다. 이 값을 결과에 누적하면 해당 칸에서 끝나는 모든 크기의 정사각형 개수까지 함께 더해지게 됩니다.

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

• 행렬이 [[1]]처럼 1 하나뿐이라면 1을 반환합니다.
• rows := 행렬의 행(row) 개수
• cols := 행렬[0]의 열(column) 개수
• result := 0으로 초기화
• row를 0부터 rows-1까지 반복합니다.
  • col을 0부터 cols-1까지 반복합니다.
    • row가 0이거나 col이 0인 경우(첫 번째 행 또는 첫 번째 열):
      • matrix[row][col]이 1이면 result를 1 증가시킵니다.
    • 그 외의 경우 matrix[row][col]이 1이라면:
      • square := 1 + min(matrix[row-1][col], matrix[row][col-1], matrix[row-1][col-1])
      • matrix[row][col] := square로 갱신
      • result := result + square

파이썬 구현 예제

def solve(matrix):
    if matrix == [[1]]:
        return 1
    rows = len(matrix)
    cols = len(matrix[0])
    result = 0
    for row in range(rows):
        for col in range(cols):
            if (row == 0 or col == 0):
                if matrix[row][col] == 1:
                    result += 1
            elif matrix[row][col] == 1:
                square = min(matrix[row-1][col], min(matrix[row][col-1], matrix[row-1][col-1])) + 1
                matrix[row][col] = square
                result += square
    return result

matrix = [[0,1,1,1],[1,1,1,1],[0,1,1,1]]
print(solve(matrix))

입력

[[0,1,1,1],[1,1,1,1],[0,1,1,1]]

출력

15

동작 원리 설명

첫 번째 행과 첫 번째 열은 비교 기준이 되는 이전 칸이 없으므로, 해당 칸의 값이 1이면 길이 1짜리 정사각형 하나만 존재합니다. 나머지 칸에서는 왼쪽, 위쪽, 대각선 방향의 세 값 중 최솟값에 1을 더함으로써 현재 칸에서 만들 수 있는 최대 정사각형의 크기를 구하고, 이를 결과에 누적합니다.

이 방식은 행렬의 모든 칸을 한 번씩만 방문하면 되므로 시간 복잡도는 O(m × n)이며, 추가 공간 없이 입력 행렬을 그대로 활용할 수 있어 매우 효율적입니다.