Python numpy中的块标量向量乘法

我有一个大的一维数组

x

，它是由不同长度的较小数组

x_0，…，x_m-1

串联而成的。我还知道每个

xi

长度的列表

L

。给定长度为

m

的数组

a

，目标是计算平面数组

[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中，有没有比这个简单的实现更简单（更快、更具Python风格等）的方法

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。

用

np重复
a
的元素。重复
并执行元素乘法-

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

不合适的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

@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

就地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

@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版本需要更长的时间（编译开销）。
这似乎不是更快，而是简单得多。@cheyp如果每个
xi
都以百万为单位，那么循环执行可能更好，因为重复
会创建类似的副本。酷，我有类似的基准测试。在我将numba与numpy结合之前，但看起来普通的numba（与
fastmath
）要好得多。