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

from typing import List, Union
import numpy as np

def calculate_matrix_mean(matrix: List[List[Union[int, float]]], mode: str) -> List[float]:
    m = len(matrix)
    n = len(matrix[0])
    sum = []
    if mode == "row":
        for i in range(m):
            sum1 = 0
            for j in range(n):
                sum1 += matrix[i][j]
            sum1 = sum1/n
            sum.append(sum1) 
    else:
        for j in range(n):
            sum1 = 0
            for i in range(m):
                sum1 += matrix[i][j]
            sum1 = sum1/m
            sum.append(sum1)       
    return sum
    
def main():
    matrix = eval(input())
    mode = input()
    result = calculate_matrix_mean(matrix, mode)
    print(result)

if __name__ == "__main__":
    main()