Skip to content
Advertisement

How do I reduce the use of for loops using numpy?

Basically, I have three arrays that I multiply with values from 0 to 2, expanding the number of rows to the number of products (the values to be multiplied are the same for each array). From there, I want to calculate the product of every combination of rows from all three arrays. So I have three arrays

A = np.array([1, 2, 3])
B = np.array([1, 2, 3])
C = np.array([1, 2, 3])

and I’m trying to reduce the operation given below

search_range = np.linspace(0, 2, 11)
results = np.array([[0, 0, 0]])

for i in search_range:
    for j in search_range:
        for k in search_range:
            sm = i*A + j*B + k*C
            results = np.append(results, [sm], axis=0)

What I tried doing:

A = np.array([[1, 2, 3]])
B = np.array([[1, 2, 3]])
C = np.array([[1, 2, 3]])

n = 11
scale = np.linspace(0, 2, n).reshape(-1, 1)
A = np.repeat(A, n, axis=0) * scale
B = np.repeat(B, n, axis=0) * scale
C = np.repeat(C, n, axis=0) * scale

results = np.array([[0, 0, 0]])

for i in range(n):
    A_i = A[i]
    for j in range(n):
        B_j = B[j]
        C_k = C
        sm = A_i + B_j + C_k
        results = np.append(results, sm, axis=0)

which only removes the last for loop. How do I reduce the other for loops?

Advertisement

Answer

You can get the same result like this:

search_range = np.linspace(0, 2, 11)
search_range = np.array(np.meshgrid(search_range, search_range, search_range))
search_range = search_range.T.reshape(-1, 3)

sm = search_range[:, 0, None]*A + search_range[:, 1, None]*B + search_range[:, 2, None]*C
results = np.concatenate(([[0, 0, 0]], sm))

Instead of using three nested loops to get every combination of elements in the “search_range” array, I used the meshgrid function to convert “search_range” to a 2D array of every possible combination and then instead of i, j and k you can use the 3 items in the arrays in the “search_range”.

And finally, as suggested by @Mercury you can use indexing for the new “search_range” array to generate the result. For example search_range[:, 1, None] is an array in shape of (1331, 1), containing singleton arrays of every element at index of 0 in arrays in the “search_range”. That concatenate is only there because you wanted the results array to have default value of [[0, 0, 0]], so I appended sm to it; Otherwise, the sm array contains the answer.

User contributions licensed under: CC BY-SA
5 People found this is helpful
Advertisement