如何修正数组分组代码，实现指定数量的近似等长分组？

我有一个包含302个元素的数组，想要将其拆分为n个近似等长的子组（示例中n=6时的分组结果如下）。当前代码在n=6时能正常工作，但设置n=51时失效，生成了60个分组。请问如何修改代码实现正确分组？

原代码及示例输出 #

import numpy as np

n = 6
group_num = np.arange(302) // (302 // n)
group_num[group_num == n] = n - 1
# 输出结果
array([
    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
    0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
    1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
    1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
    2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
    2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3,
    3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
    3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
    3, 3, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4,
    4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4,
    4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,
    5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,
    5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
    dtype=int32
)

问题原因 #

原代码用302 // n计算每组基础长度，但当总元素数无法被分组数整除时，这个逻辑会导致分组数量超出n。比如n=51时，302//51=5，302//5=60，所以会生成0到59共60个分组，后续的替换操作group_num[group_num ==n] =n-1完全起不到作用——此时最大分组号是59，远大于51。

解决方案 #

方法一：使用numpy自带的array_split函数 #

numpy的array_split函数可以直接将数组拆分为指定数量的近似等长子数组，内部已经处理了余数分配问题，无需手动计算分组号：

import numpy as np

arr = np.arange(302)
n = 51
groups = np.array_split(arr, n)
# 验证分组数量
print(len(groups))  # 输出51
# 查看每组长度：前47组长度为6，剩下4组长度为5（302=51*5+47）
for i, g in enumerate(groups[:5]):
    print(f"分组{i}长度：{len(g)}")

方法二：手动生成正确的分组号数组 #

如果需要手动生成分组号数组，可通过计算每个元素对应的分组索引，确保分组数量严格为n：

import numpy as np

total = 302
n = 51
remainder = total % n
base_len = total // n

# 生成分组号：前remainder个分组各占base_len+1个元素，剩余分组各占base_len个元素
group_num = np.concatenate([
    np.full(base_len + 1, i) for i in range(remainder)
] + [
    np.full(base_len, i) for i in range(remainder, n)
])

print(len(np.unique(group_num)))  # 输出51
print(group_num)

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