牛客题解 | 重塑矩阵

题目

题目链接

矩阵重塑是将一个矩阵转换为另一个形状的过程，前提是新形状的元素总数与原矩阵相同。数学表达式为：
如果原矩阵 A 的形状为 (m, n)，则重塑后的矩阵 B 的形状为 (p, q)，需要满足：
$$p\times q=m\times n$$

重塑操作可以通过以下代码实现：

def reshape_matrix(a: List[List[Union[int, float]]], new_shape: Tuple[int, int]) -> List[List[Union[int, float]]]:
    flat_list = [item for sublist in a for item in sublist]  # 将矩阵展平
    p, q = new_shape
    if p * q != len(flat_list):
        return -1
    return [flat_list[i * q:(i + 1) * q] for i in range(p)]  # 重塑为新形状

当然也可以使用numpy库的reshape方法简化计算

def reshape_matrix(a: List[List[Union[int, float]]], new_shape: Tuple[int, int]) -> List[List[Union[int, float]]]:
    import numpy as np
    if len(a) * len(a[0]) != new_shape[0] * new_shape[1]:
        return -1
    return np.array(a).reshape(new_shape).tolist()