NumPy：生成包含特定元素的随机子集的最快方法

Question

我需要生成一个包含索引为i元素的大小为k的子集，但要求这个子集是从大小为n的集合中随机选取的。简而言之，就是这样的：

import numpy as np

n = 7
k = 4
i = 2

subset = np.random.choice(n, k, replace=False)
while i not in subset:
    subset = np.random.choice(n, k, replace=False)
print(subset)

但也许有更快的方法吗？变量n和k的值相当小（如10或20），但我需要对不同的值进行多次这种抽样，所以最好有一些快速的方法。

- Yauhen Yakimenka

---

随机子集中元素的顺序也需要是随机的吗？ - Ben

好的，如果没有的话，可以通过 np.random.shuffle 进行修复。 - Yauhen Yakimenka

这与从一个无限制的n-1集生成k-1子集，然后再添加元素i有何不同？ - Seb

这并不是我在寻找更快的东西。 - Yauhen Yakimenka

Answer 1

方法 #1

问题在于必须使用i，而其余部分必须是随机的，但是要均匀分布。因此，我们可以对其余部分使用np.random.choice，并将i添加到末尾，然后进行随机化处理，这就是我们的输出。此外，不需要迭代。

因此，结果如下 -

def create_rand_ar(n,k,i):
    sel_ar = np.r_[:i,i+1:n]
    sel_ar_incl_i = np.r_[i,np.random.choice(sel_ar, k-1, replace=False)]
    np.random.shuffle(sel_ar_incl_i) # skip if order does not matter
    return sel_ar_incl_i

为了验证我们始终有 i 在那里，而其余部分具有相等的可能性在输出中，请在大量迭代上运行并检查出现次数的计数，应该是均匀的-

In [84]: n = 7
    ...: k = 4
    ...: i = 2

In [85]: outputs = np.array([create_rand_ar(n,k,i) for _ in range(10000)])

In [87]: np.bincount(outputs.ravel())
Out[87]: array([ 5023,  5061, 10000,  4992,  4902,  5006,  5016])

方法二

另一种方法是在[0,1)范围内创建一个均匀随机数组，将第i个元素设置为某个值< 0。然后，进行高效的argpartition并选择前k个元素，这保证了包括i且这些元素为最小值。 因此，这就是我们的输出值。

def create_rand_ar_v2(n,k,i):
    r = np.random.rand(n)
    r[i] = -1
    return r.argpartition(k)[:k]

验证分发 -

In [168]: outputs = np.array([create_rand_ar_v2(n,k,i) for _ in range(10000)])

In [169]: np.bincount(outputs.ravel())
Out[169]: array([ 4946,  5055, 10000,  5071,  4972,  5038,  4918])

时间 -

In [165]: n = 7
     ...: k = 4
     ...: i = 2

In [166]: %timeit create_rand_ar(n,k,i)
10000 loops, best of 3: 107 µs per loop

In [167]: %timeit create_rand_ar_v2(n,k,i)
100000 loops, best of 3: 2.27 µs per loop

- Divakar

---

请注意，方法2会给出一个随机子集，但是该子集的顺序并不完全随机。 - Ben

@Ben 是的，这是由于argpartition的工作方式固有的，顺序不是完全随机的，也因为我们将i添加为最低元素。因此，如果要保持顺序随机，我们可以在那里添加一个额外的np.random.shuffle()层。 - Divakar

它使版本2变慢了15倍，但仍然更快。 - Yauhen Yakimenka

Answer 2

很好的问题，一个随机选择，总是包含一个值i。 一种方法是先删除i，然后洗牌，再替换i，确保它始终存在，但从不超过一次。我假设您的n集合是可以索引的其他数据：

i = 3
k = 7
n = len(data)

idx = list(range(n))
idx.remove(i)
idx = np.append(np.random.choice(idx, k - 1, replace=False), i)
subset = data[idx]

- daveg

Answer 3

也许不是最快的，但我认为它是最优雅的竞争者：

subset = np.random.choice(n, k, replace=False)
ipos = np.random.randint(k)
subset = (subset + (i - subset[ipos])) % n
subset
# array([0, 6, 4, 2])

- Paul Panzer

Answer 4

来晚了，但这个怎么样？

elems = np.array(['a', 'b', 'c', 'd', 'e', 'f'])
guaranteed = 'c'
unguaranteed = elems[elems != guaranteed]

temp = np.random.randn(10, len(unguaranteed))  # 10 random subsets
temp = np.argsort(temp)[:, :k]

result = unguaranteed[temp]
result[np.arange(result.shape[0]), np.random.randint(low = 0, high = k, size = result.shape[0])] = guaranteed

- Ben

---

嗯...但是这种方法会比原来的（“生成直到i在子集中”）方法更快吗？ - Yauhen Yakimenka

你需要多少样本？10个，100个，1000个，1M个？这是运行一些计时的重要变量。 - Ben

10K和100K之间。 - Yauhen Yakimenka