Question

我有一个较大的一维数组x，它是通过串联不同长度的较小数组x_0 ,..., x_m-1来获得的。我也知道每个L的长度列表x_i。给定长度为a的数组m，目标是计算平面数组[a[0]*x0, a[1]*x1,...]。

例如，如果我有x = np.array([1,2,3,4,5])和a=np.array([2,-1]), L = [2,3]，则结果应为 np.array([2,4,-3,-4,-5])

在numpy中，有没有比这种天真的实现方法更简单（更快，更Pythonic等）的方法？

L.insert(0,0)
cs = np.cumsum(L)
y = np.empty(x.shape) 
for i in range(m):
    y[cs[i]:cs[i+1]] = a[i] * x[cs[i]:cs[i+1]]

我也可以在Numba做到这一点。

m大约有数百个，每个x_i的长度约为1e6。

Answer 1

用a重复np.repeat的元素并执行逐元素乘法-

y = x*np.repeat(a,L)

Answer 2

版本不正确的Numba版本

@nb.njit(fastmath=True)
def mult(x,a,L):
  out=np.empty(x.shape[0],dtype=x.dtype)
  ii=0
  for i in range(L.shape[0]):
    for j in range(L[i]):
      out[ii]=x[ii]*a[i]
      ii+=1
  return out

就地Numba版本

@nb.njit(fastmath=True)
def mult(x,a,L):
  ii=0
  for i in range(L.shape[0]):
    for j in range(L[i]):
      x[ii]=x[ii]*a[i]
      ii+=1
  return x

时间

L=np.random.randint(low=1000,high=2000,size=500)
x=np.random.rand(L.sum())
a=np.random.rand(L.shape[0])

Divakar's version:          6.4ms
Out of place Numba version: 2.8ms
In place Numba version:     1.2ms

请注意，首次调用Numba版本会花费更长的时间（编译开销）。

numpy中的块标量向量乘法

2 个答案: