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

파이썬으로 정사각 행렬의 대각선 합 구하기

정사각 행렬(square matrix)이 하나 주어졌다고 가정해 봅시다. 우리가 구해야 할 값은 이 행렬의 대각선 요소들의 총합입니다. 즉, 주대각선(primary diagonal)부대각선(secondary diagonal)에 있는 모든 요소를 더하되, 두 대각선이 교차하는 중앙 요소는 중복 계산되지 않도록 한 번만 포함해야 합니다.

예를 들어 다음과 같은 4×4 행렬이 입력으로 주어진 경우를 살펴보겠습니다.

10596
81532
38123
21173

주대각선 요소는 [10, 15, 12, 3]으로 그 합은 40이며, 부대각선 요소는 [6, 3, 8, 2]로 그 합은 19입니다. 따라서 최종 결과는 40 + 19 = 59가 됩니다.

문제 해결 접근 방법

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

  • m := 행렬의 행(row) 개수

  • m이 1이라면, 단일 요소 행렬이므로 matrix[0][0]을 바로 반환합니다.

  • count := 0 으로 초기화합니다.

  • i를 0부터 m-1까지 반복하면서 다음을 수행합니다.

    • count := count + matrix[i][i]  (주대각선 요소 누적)

    • count := count + matrix[i][-1-i]  (부대각선 요소 누적)

  • m이 홀수라면, 두 대각선이 교차하는 중앙 요소가 두 번 더해졌으므로 한 번 빼줍니다.

    • ind := m / 2의 몫

    • count := count - matrix[ind][ind]

  • count를 반환합니다.

이 알고리즘은 행렬을 한 번만 순회하므로 시간 복잡도는 O(n)이며, 추가 공간 없이 제자리에서 계산할 수 있어 매우 효율적입니다.

파이썬 구현 예제

아래 구현을 통해 더 잘 이해해 보겠습니다.

def solve(matrix):
    m = len(matrix)
    if m == 1: return matrix[0][0]

    count = 0
    for i in range(m):
        count += matrix[i][i]
        count += matrix[i][-1 - i]

    if m % 2 == 1: count -= matrix[m // 2][m // 2]

    return count

matrix = [[10,5,9,6],[8,15,3,2],[3,8,12,3],[2,11,7,3]]
print(solve(matrix))

입력

[[10,5,9,6],[8,15,3,2],[3,8,12,3],[2,11,7,3]]

출력

59

여기서 파이썬의 음수 인덱싱 기능인 matrix[i][-1 - i]를 활용하면 별도의 인덱스 계산 없이 간결하게 부대각선 요소에 접근할 수 있다는 점이 특징입니다. 또한 행렬 크기가 홀수일 때 중앙 요소를 한 번 차감하는 처리만으로 교차 요소 중복 문제를 손쉽게 해결할 수 있습니다.