为numpy的array_split函数添加随机化的通用实现方案问询

我们可以使用np.array_split快速将数组按指定份数分割，同时保留元素原有顺序：

常规分割示例 #

可整除情况 #

当数组长度能被分割份数整除时，每份元素数量均等：

import numpy as np
arr = np.array([0,1,2,3,4,5,6,7,8])
pieces = 3
np.array_split(arr, pieces)
# 输出：[array([0, 1, 2]), array([3, 4, 5]), array([6, 7, 8])]

不可整除情况 #

当数组长度无法被分割份数整除时，np.array_split默认让前若干份多1个元素，最后一份元素数量更少：

arr = np.array([0,1,2,3,4,5,6,7])
pieces = 3
np.array_split(arr, pieces)
# 输出：[array([0, 1, 2]), array([3, 4, 5]), array([6, 7])]

需求：实现均等概率的随机化分割 #

我们需要让所有合法的不均等分割结果以相等概率出现。比如上述8元素分3份的场景，以下三种结果应等概率输出：

# 结果1
[array([0, 1]), array([2, 3, 4]), array([5, 6, 7])]
# 结果2
[array([0, 1, 2]), array([3, 4]), array([5, 6, 7])]
# 结果3
[array([0, 1, 2]), array([3, 4, 5]), array([6, 7])]

通用解决方案 #

核心思路：先计算基础每份长度和需要多分配1个元素的份数，再随机打乱这些长度的分配顺序，最后按随机后的长度分割数组。

实现代码：

import numpy as np

def random_array_split(arr, pieces):
    arr_size = arr.size
    base_len = arr_size // pieces
    remainder = arr_size % pieces
    
    # 构造分割长度列表：remainder份为base_len+1，剩余为base_len
    segment_lengths = [base_len + 1] * remainder + [base_len] * (pieces - remainder)
    # 随机打乱长度顺序
    np.random.shuffle(segment_lengths)
    
    # 计算分割点并执行分割
    split_points = np.cumsum(segment_lengths)[:-1]
    return np.split(arr, split_points)

测试示例 #

示例1：8元素分3份

arr = np.array([0,1,2,3,4,5,6,7])
pieces = 3
# 多次调用会等概率返回三种合法分割结果
print(random_array_split(arr, pieces))

示例2：10元素分6份

arr = np.array([0,1,2,3,4,5,6,7,8,9])
pieces = 6
# 输出包含4个长度为2、2个长度为1的子数组，顺序随机
print(random_array_split(arr, pieces))

方案优势 #

• 适配所有数组长度与分割份数的组合
• 所有合法分割结果出现概率完全均等
• 严格保留原数组的元素顺序
• 基于Numpy原生函数实现，执行效率高

内容的提问来源于stack exchange，提问作者SerKo