Python으로 행렬의 i번째 행과 i번째 열의 합이 같은지 확인하는 방법

문제 개요

2차원 행렬(매트릭스)이 주어졌을 때, i번째 행의 합과 i번째 열의 합이 서로 같은지 확인하는 프로그램을 작성해야 합니다.

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

이 경우 출력 결과는 True입니다. 첫 번째 행의 합은 (2 + 3 + 4 + 5) = 14이고, 첫 번째 열의 합은 (2 + 10 + 1 + 1) = 14로 두 값이 동일하기 때문입니다.

해결 접근 방법

이 문제는 다음 단계를 통해 해결할 수 있습니다.

• 행렬의 행(row) 개수와 열(column) 개수를 구합니다.
• 각 인덱스 i에 대해 다음을 반복합니다.
  • 행 합계(total_row)와 열 합계(total_col)를 0으로 초기화합니다.
  • j를 0부터 열 개수 - 1까지 반복하면서 total_row에는 mat[i][j]를 더하고, total_col에는 mat[j][i]를 더합니다.
  • 반복이 끝난 후 total_row와 total_col이 같으면 True를 반환합니다.
• 모든 인덱스를 검사한 후에도 일치하는 경우가 없다면 False를 반환합니다.

예제 코드

def solve(mat):
    row = len(mat)
    col = len(mat[0])
    
    for i in range(row):
        total_row = 0
        total_col = 0
        for j in range(col):
            total_row += mat[i][j]
            total_col += mat[j][i]
        
        if total_row == total_col:
            return True
    
    return False

matrix = [
    [2,3,4,5],
    [10,6,4,2],
    [1,4,6,7],
    [1,5,6,7]
]

print(solve(matrix))

입력 예시

[
[1,2,3,4],
[9,5,3,1],
[0,3,5,6],
[0,4,5,6]
]

출력 결과

True

코드 설명 및 시간 복잡도

위 코드는 이중 반복문을 사용하여 각 인덱스 i에 대해 i번째 행의 모든 요소와 i번째 열의 모든 요소를 동시에 더합니다. 내부 반복문이 한 번 완료될 때마다 행 합계와 열 합계를 비교하여 일치하면 즉시 True를 반환함으로써 불필요한 연산을 줄일 수 있습니다.

이 알고리즘의 시간 복잡도는 O(n²)입니다. 여기서 n은 행렬의 크기이며, n×n 정방 행렬을 기준으로 할 때 각 인덱스마다 전체 행과 열을 순회해야 하기 때문입니다. 공간 복잡도는 추가 변수만 사용하므로 O(1)입니다.