Numpy插入行以匹配最大行数

import numpy as np

def insert_rows_to_match_max_rows(array, rows_to_insert, max_rows):
    # 获取数组的当前行数
    current_rows = array.shape[0]
    
    if current_rows < max_rows:
        # 计算需要插入的行数
        rows_needed = max_rows - current_rows
        
        # 创建待插入的数据
        data_to_insert = np.zeros((rows_needed, array.shape[1]))
        
        # 将原数组和待插入的数据堆叠在一起
        new_array = np.vstack((array, data_to_insert))
    else:
        # 删除数组的前几行，使得数组的行数等于最大行数
        new_array = np.delete(array, np.s_[:current_rows-max_rows+rows_to_insert], axis=0)
        
        # 将删除后的数组和待插入的数据堆叠在一起
        new_array = np.vstack((new_array, rows_to_insert))
    
    return new_array