在具有不同多重性但相同维度的数组上同时使用 numpy 重复

Question

我有两个长度相同的 trival 数组，tmp_reds 和 tmp_blues：

npts = 4
tmp_reds = np.array(['red', 'red', 'red', 'red'])
tmp_blues = np.array(['blue', 'blue', 'blue', 'blue'])

我正在使用 np.repeat 来创建多重性：

red_occupations = [1, 0, 1, 2]
blue_occupations = [0, 2, 0, 1]

x = np.repeat(tmp_reds, red_occupations)
y = np.repeat(tmp_blues, blue_occupations)

print(x)
['red' 'red' 'red' 'red']

print(y)
['blue' 'blue' 'blue']

我想要的是以下 x 和 y 的组合：

desired_array = np.array(['red', 'blue', 'blue', 'red', 'red', 'red', 'blue'])

所以，desired_array的定义方式如下：

(1) 应用 red_occupations 的第一个元素的多重性

(2) 应用 blue_occupations 的第一个元素的多重性

(3) 应用 red_occupations 的第二个元素的多重性

(4) 应用 blue_occupations 的第二个元素的多重性

...

(2*npts-1) 应用 red_occupations 的 npts 元素的多重性

(2*npts) 应用了 blue_occupations 的 npts 元素的多重性

所以这似乎是对 np.repeat 正常用法的简单概括。通常， np.repeat 完全执行上述操作，但使用单个数组。有谁知道一些巧妙的方法来使用多维数组然后被展平，或者其他类似的技巧，可以用 np.repeat 完成这个？

我总是可以在不使用 numpy 的情况下创建 desired_array，使用简单的压缩 for 循环和连续的列表追加。不过实际问题有npts~1e7，速度很关键。

Answer 1

对于一般情况 -

# Two 1D color arrays
tmp1 = np.array(['red', 'red', 'red', 'green'])
tmp2 = np.array(['white', 'black', 'blue', 'blue'])

# Multiplicity arrays
color1_occupations = [1, 0, 1, 2]
color2_occupations = [0, 2, 0, 1]

# Stack those two color arrays and two multiplicity arrays separately
tmp12 = np.column_stack((tmp1,tmp2))
color_occupations = np.column_stack((color1_occupations,color2_occupations))

# Use np.repeat to get stacked multiplicities for stacked color arrays
out = np.repeat(tmp12,color_occupations.ravel())

给我们 -

In [180]: out
Out[180]: 
array(['red', 'black', 'black', 'red', 'green', 'green', 'blue'], 
      dtype='|S5')