Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

파이썬으로 2진 행렬을 k개 조각으로 나누는 방법의 수 세기

이번 문제에서는 0과 1로 이루어진 2진 행렬(binary matrix)과 정수 k가 주어집니다. 우리는 이 행렬을 k개의 조각으로 나누되, 각 조각에 최소한 하나의 1이 포함되도록 해야 합니다.

단, 잘라내는 과정에는 다음과 같은 규칙이 있습니다.

  • 먼저 자를 방향을 선택합니다 — 세로(vertical) 또는 가로(horizontal).
  • 행렬에서 잘라낼 인덱스를 지정하여 두 부분으로 나눕니다.
  • 세로로 자른 경우: 왼쪽 부분은 더 이상 자를 수 없고, 오른쪽 부분만 계속 잘라야 합니다.
  • 가로로 자른 경우: 위쪽 부분은 더 이상 자를 수 없고, 아래쪽 부분만 계속 잘라야 합니다.

이러한 규칙을 지키면서 행렬을 나눌 수 있는 서로 다른 방법의 수를 구하는 것이 목표입니다. 만약 답이 매우 커질 경우에는 결과를 (10^9 + 7)로 나눈 나머지를 반환합니다.

입력 예시

예를 들어 입력이 다음과 같다고 가정해 보겠습니다.

1
1
0
1
0
1
1
1
1

k = 2일 때 출력은 4가 됩니다. 세로로 두 번 자르는 경우와 가로로 두 번 자르는 경우가 각각 가능하기 때문입니다.

해결 접근 방법

이 문제는 접미사 합(suffix sum)과 재귀(DFS)를 활용하면 효율적으로 풀 수 있습니다. 각 위치에서 시작하는 영역에 포함된 1의 개수를 미리 계산해 두면, 특정 영역을 잘라냈을 때 남은 부분에 1이 존재하는지 빠르게 판단할 수 있습니다.

구체적인 단계는 다음과 같습니다.

  • p := 10^9 + 7 (나눗셈에 사용할 모듈로 값)
  • m := 행렬의 행(row) 개수, n := 행렬의 열(column) 개수
  • counts := 비어 있는 맵(map) 생성
  • i를 m-1부터 0까지 역순으로 순회하며:
    • j를 n-1부터 0까지 역순으로 순회하며:
      • counts[i, j] := counts[i + 1, j] + counts[(i, j + 1)] − counts[(i + 1, j + 1)] + matrix[i, j]
        (즉, (i, j) 위치부터 행렬 끝까지의 1의 개수를 누적합으로 저장)
  • f(x, y, c) 함수를 정의합니다. 여기서 x, y는 현재 영역의 시작 좌표, c는 남은 자르기 횟수입니다.
    • count := counts[x, y]
    • c가 0이라면, count > 0일 때 1을 반환하고 그렇지 않으면 0을 반환합니다.
    • ans := 0으로 초기화
    • i를 x + 1부터 m − 1까지 순회하며, 0 < counts[(i, y)] < count인 경우 ans += f(i, y, c − 1)를 수행합니다. (세로로 자르는 경우)
    • j를 y + 1부터 n − 1까지 순회하며, 0 < counts[(x, j)] < count인 경우 ans += f(x, j, c − 1)를 수행합니다. (가로로 자르는 경우)
    • ans mod p를 반환합니다.
  • 메인 메서드에서 f(0, 0, k − 1)을 호출하고 그 결과를 반환합니다.

파이썬 구현 코드

아래 코드를 통해 더 자세히 이해해 보겠습니다.

from collections import defaultdict
class Solution:
    def solve(self, matrix, k):
        p = 10 ** 9 + 7

        m, n = len(matrix), len(matrix[0])
        counts = defaultdict(int)
        for i in range(m)[::-1]:
            for j in range(n)[::-1]:
                counts[(i, j)] = (counts[(i + 1, j)] + counts[(i, j + 1)] - counts[(i + 1, j + 1)] + matrix[i][j])

        def f(x, y, c):
            count = counts[(x, y)]
            if c == 0:
                return 1 if count > 0 else 0

            ans = 0
            for i in range(x + 1, m):
                if 0 < counts[(i, y)] < count:
                    ans += f(i, y, c - 1)
            for j in range(y + 1, n):
                if 0 < counts[(x, j)] < count:
                    ans += f(x, j, c - 1)

            return ans % p
        return f(0, 0, k - 1)

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

입력

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

출력

4

정리

핵심 아이디어는 접미사 합 배열을 사용해 임의의 위치에서 시작하는 하위 행렬에 포함된 1의 개수를 O(1) 시간에 확인할 수 있게 하는 것입니다. 이후 재귀 함수를 통해 가능한 모든 자르기 위치를 탐색하면서, 잘라낸 뒤 남은 영역에 반드시 1이 하나 이상 남아 있는 경우만 유효한 방법으로 카운트합니다. 이렇게 하면 불필요한 탐색을 줄여 효율적으로 전체 경우의 수를 구할 수 있습니다.