在 Numpy Python 中按一个顺序格式化两个数组的排列

Question

我正在尝试订购zeroes 和ones 以安排order。预期的输出是我想要达到的。最好不使用列表推导。

import numpy as np

order = np.array([0,1,0,1,0])
zeroes= np.array([10,55, 30])
ones = np.array([3,8])

预期输出

[10, 3, 55, 8, 30] 

Answer 1

这个怎么样（没有 Python 循环：在 200k 元素上测试时，比列表理解快 750 倍）：

# note: updated version: faster and more robust to faulty input
def altcat(zeroes, ones, order):
    i0 = np.nonzero(order == 0)[0][:len(zeroes)]
    i1 = np.nonzero(order == 1)[0][:len(ones)]
    
    z = np.zeros_like(order, dtype=zeroes.dtype)
    z[i0] = zeroes[:len(i0)]
    z[i1] = ones[:len(i1)]
    return z

关于你的例子：

>>> altcat(zeroes=np.array([10,55, 30]), ones=np.array([3,8]),
...        order=np.array([0,1,0,1,0]))
array([10,  3, 55,  8, 30])

速度

# set up
n = 200_000
np.random.seed(0)
order = np.random.randint(0, 2, size=n)
n1 = order.sum()
n0 = n - n1
ones = np.random.randint(100, size=n1)
zeroes = np.random.randint(100, size=n0)

# for comparison, a method proposed elsewhere, based on lists
def altcat_list(zeroes, ones, order):
    zeroes = list(zeroes)
    ones = list(ones)
    return [zeroes.pop(0) if i == 0 else ones.pop(0) for i in order]

测试：

a = %timeit -o altcat(zeroes, ones, order)
# 2.38 ms ± 573 ns per loop (mean ± std. dev. of 7 runs, 100 loops each)

b = %timeit -o altcat_list(zeroes, ones, order)
# 1.84 s ± 1.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

b.average / a.average
# 773.59

注意：我最初尝试使用 n = 1_000_000，但虽然 altcat 在 12.4 毫秒内完成，但基于列表的版本将花费很长时间，我不得不停止它。

似乎基于列表的方法比O(n)差（100K：0.4s；200K：1.84s；400K：10.4s）。

附录

如果你真的想用列表理解而不是纯numpy，那么至少考虑一下：

def altcat_list_mod(zeroes, ones, order):
    it = [iter(zeroes), iter(ones)]
    return [next(it[i]) for i in order]

这比 altcat_list() 快，但仍然比 altcat() 慢近 25 倍：

# on 200k elements

c = %timeit -o altcat_list_mod(zeroes, ones, order)
# 60 ms ± 24.3 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

c.average / a.average
# 24.93