如何根据列的总和来组织二维数组的列？

def organize_array_by_column_sum(column_sums):
    num_columns = len(column_sums)  # 确定列数
    num_rows = max(column_sums)  # 确定行数，选择最大的列总和作为行数

    array = []
    for i in range(num_columns):
        column_sum = column_sums[i]  # 获取当前列的总和
        column = [i + 1] * column_sum  # 根据总和创建一维数组
        column += [0] * (num_rows - column_sum)  # 补齐长度
        array.append(column)  # 添加到二维数组中

    return array

column_sums = [3, 2, 4, 1]  # 列总和

result = organize_array_by_column_sum(column_sums)
print(result)

[[1, 1, 1, 0], [2, 2, 0, 0], [3, 3, 3, 3], [4, 0, 0, 0]]