题解 | 重塑矩阵

from re import A
from typing import List, Tuple, Union
import numpy as np


def reshape_matrix(
    a: List[List[Union[int, float]]], new_shape: Tuple[int, int]
) -> List[List[Union[int, float]]]:
    flattened = [num for row in a for num in row]
    r, c = new_shape
    # 检查元素总数是否匹配
    if r * c != len(flattened):
        return -1
    # 生成新的矩阵
    return [flattened[i * c : (i + 1) * c] for i in range(r)]


def main():
    try:
        a = eval(input())
        new_shape = eval(input())
        result = reshape_matrix(a, new_shape)
        print(result)
    except Exception as e:
        print(f"输入格式错误: {e}")


if __name__ == "__main__":
    main()