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

from typing import List, Union
import numpy as np

def calculate_matrix_mean(matrix: List[List[Union[int, float]]], mode: str) -> List[float]:
    np_matrix = np.array(matrix)

    if mode == 'row':

        averages = np.mean(np_matrix, axis=1).tolist()

    else:

        averages = np.mean(np_matrix, axis=0).tolist()

    return [float(avg) for avg in averages]
    
def main():
    matrix = eval(input())
    mode = input()
    result = calculate_matrix_mean(matrix, mode)
    print(result)

if __name__ == "__main__":
    main()