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 *.

• python
• numpy

• What is indices? What are the constraints on i/j/k? – mozway Commented Jan 17 at 19:28
• Indices is a list of tuples of integers. No constraints on i/j/k, but the list of indices will be known ahead of time. – Will Commented Jan 17 at 19:30

1 Answer 1

You 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