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

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 == 'row':
        return np.mean(matrix,axis=1).tolist()
    elif mode == 'column':
        return np.mean(matrix,axis=0).tolist()
    else:
        pass
    
def main():
    matrix = eval(input())
    mode = input()
    result = calculate_matrix_mean(matrix, mode)
    print(result)

if __name__ == "__main__":
    main()