# 행렬 대각합

import numpy as np

# 행렬
a = np.array([[1, 2, 3], [3, 4, 5], [4, 5, 6]])

# 아인슈타인 표기법을 사용한 대각합 계산
trace = np.einsum( 'ii->', a )
# 'ii->'는 동일한 인덱스의 대각 요소들을 합산하라는 의미입니다.

print( "a = ", a, sep="\n" )
print("Trace of the matrix:", trace)
# Output: 5 (1 + 4)