题解 | 按行或列计算平均值

from typing import List, Union
import numpy as np


def calculate_matrix_mean(matrix: List[List[Union[int, float]]], mode: str) -> List[float]:
    matrix = np.array(matrix)
    if mode == 'column':
        matrix = np.transpose(matrix)
    result = []
    for m in matrix:
        result.append(np.mean(m))
    return result


def main():
    matrix = eval(input())
    mode = input()
    result = calculate_matrix_mean(matrix, mode)
    print(result)


if __name__ == "__main__":
    main()