CODEWITHSUNDEEP
pythonnumpy
slaw
slaw

Reputation: 6869

Convert Python For-Loop to NumPy Operations

I have a NumPy array full of indices:

size = 100000
idx = np.random.randint(0, size, size=size)

And I have a simple function that loops over the indices and does:

out = np.zeros(size, dtype=np.int)

for i in range(size):
    j = idx[i]
    out[min(i, j)] = out[min(i, j)] + 1
    out[max(i, j)] = out[max(i, j)] - 1

return np.cumsum(out)

This is quite slow when size is large and I am hoping to find a faster way to accomplish this. I've tried this but it isn't quite right:

out = np.zeros(size, dtype=np.int)
i = np.arange(size)
j = idx[i]
mini = np.minimum(i, j)
maxi = np.maximum(i, j)

out[mini] = out[mini] + 1
out[maxi] = out[maxi] - 1

return np.cumsum(out)

Upvotes: 1

Views: 388

Answers (2)

Divakar
Divakar

Reputation: 221574

We can make use of np.bincount -

R = np.arange(size)
out = np.bincount(np.minimum(R,idx),minlength=size)
out -= np.bincount(np.maximum(R,idx),minlength=size)
final_out = out.cumsum()

Timings -

All posted solutions use cumsum at the end. So, let's time these skipping that last step -

In [25]: np.random.seed(0)
    ...: size = 100000
    ...: idx = np.random.randint(0, size, size=size)

# From this post
In [27]: %%timeit
    ...: R = np.arange(size)
    ...: out = np.bincount(np.minimum(R,idx),minlength=size)
    ...: out -= np.bincount(np.maximum(R,idx),minlength=size)
1000 loops, best of 3: 643 µs per loop

# @slaw's solution
In [28]: %%timeit
    ...: i = np.arange(size)
    ...: j = idx[i]
    ...: mini = np.minimum(i, j)
    ...: maxi = np.maximum(i, j)
    ...: 
    ...: unique_mini, mini_counts = np.unique(mini, return_counts=True)
    ...: unique_maxi, maxi_counts = np.unique(maxi, return_counts=True)
    ...: 
    ...: out = np.zeros(size, dtype=np.int)
    ...: out[unique_mini] = out[unique_mini] + mini_counts
    ...: out[unique_maxi] = out[unique_maxi] - maxi_counts
100 loops, best of 3: 13.3 ms per loop

# Loopy one from question
In [29]: %%timeit
    ...: out = np.zeros(size, dtype=np.int)
    ...: 
    ...: for i in range(size):
    ...:     j = idx[i]
    ...:     out[min(i, j)] = out[min(i, j)] + 1
    ...:     out[max(i, j)] = out[max(i, j)] - 1
10 loops, best of 3: 141 ms per loop

Upvotes: 1

slaw
slaw

Reputation: 6869

This seems to give the same answer as the for-loop

i = np.arange(size)
j = idx[i]
mini = np.minimum(i, j)
maxi = np.maximum(i, j)

unique_mini, mini_counts = np.unique(mini, return_counts=True)
unique_maxi, maxi_counts = np.unique(maxi, return_counts=True)

out = np.zeros(size, dtype=np.int)
out[unique_mini] = out[unique_mini] + mini_counts
out[unique_maxi] = out[unique_maxi] - maxi_counts

return np.cumsum(out)

Upvotes: 1

Related Questions