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

2023年1月6日 273次阅读

我有两个相同长度的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可以实现扁平化或其他类似技巧？

我总是可以创建desired_array而不使用numpy,使用简单的zipped for循环和连续列表追加.但是,实际问题有npts~1e7,速度很关键.

最佳答案 对于一般情况 –

# 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')
点赞