파이썬으로 행렬에서 조건을 만족하는 특수 요소 개수 찾기

문제 이해하기

0과 1로만 이루어진 이진 행렬(binary matrix)이 주어졌을 때, 다음 두 가지 조건을 동시에 만족하는 요소의 개수를 찾아야 합니다.

• matrix[r][c] = 1
• 같은 행(r)의 다른 모든 열 j(j ≠ c)에 대해 matrix[r][j] = 0이고, 같은 열(c)의 다른 모든 행 i(i ≠ r)에 대해 matrix[i][c] = 0

쉽게 말해, 자기 자신만 1이고 속한 행과 열의 나머지 값은 모두 0인 '특수한 위치'를 세는 문제입니다.

예제 살펴보기

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

이 경우 출력은 3입니다. 셀 (0,2), (1,0), (2,1)이 각각 해당 행과 열에서 유일한 1이므로 조건을 만족하기 때문입니다.

해결 접근 방법

핵심 아이디어는 간단합니다. 각 행과 열에 포함된 1의 개수를 미리 계산해 두면, 어떤 셀이 조건을 만족하는지 상수 시간에 판별할 수 있습니다. 알고리즘은 다음과 같습니다.

1. 행렬이 비어 있으면 0을 반환합니다.
2. 각 행의 합(row)과 각 열의 합(col)을 리스트로 계산합니다.
3. 행렬의 행 개수 m과 열 개수 n을 구합니다.
4. 결과값 res를 0으로 초기화합니다.
5. 모든 셀 (r, c)를 순회하면서 matrix[r][c]가 1이고 row[r]이 1이며 col[c]도 1인 경우 res를 1 증가시킵니다.
6. 순회가 끝나면 res를 반환합니다.

row[r]이 1이라는 것은 r번째 행 전체에서 1이 딱 하나뿐이라는 의미이고, col[c]가 1이라는 것은 c번째 열에서도 마찬가지라는 뜻입니다. 따라서 matrix[r][c]가 1이면서 이 두 조건을 함께 만족한다면 그 셀은 정확히 우리가 찾는 '특수한 요소'입니다.

파이썬 구현 코드

def solve(matrix):
    if not matrix:
        return 0

    # 각 행의 합 계산
    row = [sum(r) for r in matrix]
    # 각 열의 합 계산
    col = [sum(c) for c in zip(*matrix)]

    m, n = len(matrix), len(matrix[0])
    res = 0
    for r in range(m):
        for c in range(n):
            if matrix[r][c] == 1 and row[r] == 1 and col[c] == 1:
                res += 1
    return res

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

입력

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

출력

3

복잡도 분석

행렬의 크기를 m × n이라고 할 때, 각 행과 열의 합을 계산하는 데 O(m × n)의 시간이 걸리고, 전체 셀을 한 번 순회하는 데도 O(m × n)이 소요됩니다. 따라서 전체 시간 복잡도는 O(m × n)입니다. 추가로 사용되는 공간은 행 합과 열 합을 저장하는 리스트 두 개뿐이므로 공간 복잡도는 O(m + n)으로 매우 효율적입니다.