python - 使用条件连接 NumPy 数组

Question

我正在尝试根据一个条件连接两个 NumPy 数组，因此连接的顺序取决于该条件。

这是我现在拥有的:

a = empty([2048, 1536])
    
for i in range(2048):
    if s[i]: # s[] -> array with shape (2048, 1) with ones and zeros
        a[i] = concatenate([b[i], c[i]]) # b[] and c[] -> arrays with shape (2048, 768)
        continue
    a[i] = concatenate([c[i], b[i]]) # order changes depending on condition

我不知道如何只进行一次操作或减少计算时间。我希望它尽可能优化，但我是 NumPy 的新手，不太了解它。

谢谢，

Answer 1

我认为连接不是正确的解决方案，可以通过适当的切片来解决:

import numpy as np

dim1 = 32
dim2 = 16

a = np.empty([dim1, dim2])
b = np.ones([dim1, dim2//2])
c = np.zeros([dim1, dim2//2])
    
s =     np.ones([dim1, 1])
s[:dim1//2] = 0

for i in range(dim1):
    
    if s[i] == 0:
        a[i, :dim2//2] = b[i, :]
        a[i, dim2//2:] = c[i, :]
    else:
        
        a[i, :dim2//2] = c[i, :]
        a[i, dim2//2:] = b[i, :]