python - 在具有不同重数但相同维度的数组上同时使用 numpy repeat

Question

我有两个相同长度的平凡数组，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')