如何从numpy数组中选取符合最小连续簇长度要求的随机索引子集

思路 #

• 先将需要的总采样长度拆分为若干个不小于指定最小簇长度的随机整数，每个整数对应一个连续索引簇的长度
• 计算簇放置所需的最小空间（所有簇长度之和 + 簇之间至少1个间隔的总长度），校验原数组长度是否满足要求
• 将额外可用的空位随机分配到首段、簇间隔、尾段三个位置，保证簇的位置完全随机且无重叠
• 合并所有簇的索引得到最终结果

可运行代码 #

import numpy as np

def sample_clustered_indices(arr_len, total_sample_len, min_cluster_len=3):
    # 参数合法性校验
    if total_sample_len < min_cluster_len:
        raise ValueError(f"总采样长度不能小于最小簇长度{min_cluster_len}")
    if arr_len < total_sample_len:
        raise ValueError("原数组长度小于需要采样的总长度")

    # 1. 拆分总采样长度为若干符合最小长度要求的簇长度
    cluster_lengths = []
    remaining = total_sample_len
    while remaining > 0:
        if remaining < min_cluster_len:
            # 剩余长度不足最小簇长时，补到最后一个已生成的簇
            cluster_lengths[-1] += remaining
            break
        # 随机生成当前簇长度
        curr_len = np.random.randint(min_cluster_len, remaining + 1)
        # 避免剩余长度不足最小簇长
        if 0 < remaining - curr_len < min_cluster_len:
            continue
        cluster_lengths.append(curr_len)
        remaining -= curr_len
    # 打乱簇长度顺序，避免固定长度分布
    np.random.shuffle(cluster_lengths)
    total_clusters = len(cluster_lengths)

    # 2. 校验簇放置的最小空间要求（簇之间至少留1个未选索引避免合并为同一个簇）
    min_required_len = total_sample_len + (total_clusters - 1)
    if arr_len < min_required_len:
        raise ValueError(f"当前簇数量下，原数组长度至少需要{min_required_len}才能放下所有不重叠的簇")
    
    # 3. 随机分配额外空位，确定每个簇的位置
    extra_slots = arr_len - min_required_len
    # 将额外空位随机分配到首段、簇间隔、尾段共(total_clusters + 1)个间隙
    gaps = np.random.multinomial(extra_slots, np.ones(total_clusters + 1) / (total_clusters + 1))
    clusters = []
    curr_pos = gaps[0]
    for i in range(total_clusters):
        cl = cluster_lengths[i]
        clusters.append(np.arange(curr_pos, curr_pos + cl))
        # 跳过当前簇、1个间隔、下一个间隙的空位
        curr_pos += cl + 1 + gaps[i + 1]
    
    # 4. 合并所有簇的索引
    return np.concatenate(clusters)

测试示例 #

以你给出的长度为25的数组、采样总长度12为例：

a = np.arange(0, 25)
sampled_indices = sample_clustered_indices(arr_len=len(a), total_sample_len=12)
print(sampled_indices)
# 输出示例：[ 3  4  5  6  7  8 14 15 16 22 23 24]

方案优势 #

• 支持任意长度的输入数组，无需被最小簇长整除
• 簇长度完全随机，仅限制最小长度，可生成任意符合要求的簇长组合
• 支持自定义最小簇长度，可灵活适配其他需求

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