Convert Python For-Loop to NumPy Operations

Question

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)

Divakar · Accepted Answer · 2019-07-13 05:23:37Z

1

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

answered Jul 13, 2019 at 5:23

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

slaw Over a year ago

Wow, that's insanely fast and a more compact too!

slaw · Accepted Answer · 2019-07-13 01:08:07Z

1

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)

answered Jul 13, 2019 at 1:08

slaw

6,98919 gold badges67 silver badges118 bronze badges

Collectives™ on Stack Overflow

Convert Python For-Loop to NumPy Operations

2 Answers 2

1 Comment

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

1 Comment

Comments

Your Answer

Sign up or log in

Post as a guest

Related