如何用Numpy高效原地按指定位分割u64一维数组？

你可以通过以下两种高效方式实现需求，同时直接获取分割索引：

方式一：掩码拼接（简洁直观） #

这种方式和你原代码思路一致，同时直接计算分割索引：

import numpy as np

arr = np.arange(720, dtype=np.uint64)
np.random.shuffle(arr)

# 标记第5位为0的元素
mask = (arr & 32) == 0
# 分割索引即左半部分元素数量
split_idx = mask.sum()

# 原地更新数组：左半部分为第5位0的元素，右半部分为第5位1的元素
arr[:] = np.concatenate([arr[mask], arr[~mask]])

方式二：基于np.argpartition（低内存开销） #

如果数组规模极大，担心掩码拼接产生的临时数组占用过多内存，可以用np.argpartition实现更高效的内存利用，时间复杂度接近O(n)：

import numpy as np

arr = np.arange(720, dtype=np.uint64)
np.random.shuffle(arr)

# 生成键值：第5位为0则是0，为1则是1
key = (arr & 32) >> 5
# 计算分割索引
split_idx = np.count_nonzero(key == 0)

# 通过索引分割数组并原地更新
arr[:] = arr[np.argpartition(key, split_idx)]

注：np.argpartition会将数组中对应键值为0的元素全部移到前split_idx位置，键值为1的移到后面，内部采用类似快速排序的分治策略，无需完全排序，效率比全排序更高。

关键说明 #

• 分割索引可以直接通过mask.sum()或np.count_nonzero(key == 0)获取，无需额外计算len(left)。
• 两种方式都能实现原地修改原数组的需求，区别在于内存开销和实现复杂度：掩码拼接更易读，argpartition更适合超大规模数组。

内容的提问来源于stack exchange，提问作者Michał Šrajer