I want to compute an expression consisting of numpy arrays, but only at a given index. For example, say I have numpy arrays a
, b
, and c
and want to compute (a[i] + b[j]) * c[k, i]
. I could do something like
total = 0
for (i, j, k) in indices:
total += (a[i] * b[j]) + c[k, i]
but I'm looking for the most optimized way to do this. Is there a way to do this in pure numpy to avoid the for loop? I'm looking for a general solution that will work for arbitrary expressions involving +
and *
.
I want to compute an expression consisting of numpy arrays, but only at a given index. For example, say I have numpy arrays a
, b
, and c
and want to compute (a[i] + b[j]) * c[k, i]
. I could do something like
total = 0
for (i, j, k) in indices:
total += (a[i] * b[j]) + c[k, i]
but I'm looking for the most optimized way to do this. Is there a way to do this in pure numpy to avoid the for loop? I'm looking for a general solution that will work for arbitrary expressions involving +
and *
.
1 Answer
Reset to default 2You should be able to vectorize your operation with array indexing and np.sum
:
i, j, k = np.array(indices).T
total = np.sum((a[i] * b[j]) + c[k, i])
Example:
a = np.random.randint(0, 100, 10)
b = np.random.randint(0, 100, 20)
c = np.random.randint(0, 100, (30, 10))
indices = [(0, 19, 29), (1, 2, 3), (5, 4, 3)]
# loop
total = 0
for (i, j, k) in indices:
total += (a[i] * b[j]) + c[k, i]
# vectorial
i,j,k = np.array(indices).T
total2 = np.sum((a[i] * b[j]) + c[k, i])
assert total == total2
indices
? What are the constraints on i/j/k? – mozway Commented Jan 17 at 19:28