Sum Diagonal of Matrix in Python

def sum_diagonals(z: list[list[int]]):
    row_count = len(z)
    total_sum = 0

    for d in range(row_count):
        # Diagonal element
        total_sum += z[d][d]
        # Anti-diagonal element
        total_sum += z[row_count - 1 - d][d]

    # If matrix has odd number of rows, the middle element got
    # added twice, so subtract it once.
    if row_count % 2 != 0:
        total_sum -= z[row_count // 2][row_count // 2]

    return total_sum


def read_matrix(m: int, n: int):
    mat: list[list[int]] = []
    for r in range(m):
        row_input = input(f"Row {r:2d}: ").strip().split()
        row = list(map(int, row_input[:n]))
        mat.append(row)
    return mat


def main():
    n = int(input("Enter n: "))
    a = read_matrix(n, n)
    r = sum_diagonals(a)
    print(r)


if __name__ == "__main__":
    main()