Is there any way in numpy to vectorize the following operations?
a = np.zeros(3)
idx = np.array([1, 0, 1])
val = np.array([1, 2, 3])
for i, v in zip(idx, val):
a[i] += v
# a = [2, 4, 0]
The following code does not work as desired
a = np.zeros(3)
idx = np.array([1, 0, 1])
val = np.array([1, 2, 3])
a[idx] += val
# wrong answer: a = [2, 3, 0]
We can use np.bincount setting val as weights:
np.bincount(idx, weights=val, minlength=len(idx)).astype(int)
#array([2, 4, 0])
This is in fact doing exactly as shared in the question. By adding a weight to np.bincount, instead of:
np.bincount(idx, minlength=len(idx))
# array([1, 2, 0], dtype=int64)
Which would be the same as obtaining the counts:
for i in idx:
a[i] += 1
print(a)
# array([1., 2., 0.])
We have that out[n] += weight[i]:
for i, v in zip(idx, val):
a[i] += v
print(a)
# array([2., 4., 0.])