CODEWITHSUNDEEP
pythonnumpy
Adrian Holovaty
Adrian Holovaty

Reputation: 2419

Distance to next non-zero element in one-dimensional numpy array

I have a one-dimensional numpy array consisting of ones and zeroes, like this:

[0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1]

For each non-zero element of the array, I want to calculate the "distance" to the next non-zero element. That is, I want to answer the question "How far away is the next non-zero element?" So the result for the above array would be:

[0, 0, 0, 1, 3, 0, 0, 6, 0, 0, 0, 0, 0, 4, 0, 0, 0, 0]

Is there a built-in numpy function for this? And if not, what's the most efficient way to implement this in numpy?

Upvotes: 0

Views: 746

Answers (3)

Dani Mesejo
Dani Mesejo

Reputation: 61920

Let's try:

import numpy as np

arr = np.array([0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1])

# create output
res = np.zeros_like(arr)

# select indices non-zero
where, = np.where(arr)

# assign the indices of the non-zero the diff
res[where[:-1]] = np.diff(where)
print(res)

Output

[0 0 0 1 3 0 0 6 0 0 0 0 0 4 0 0 0 0]

Upvotes: 1

Lior Cohen
Lior Cohen

Reputation: 5745

Here is 2 liners. If you don't want override original a replace with copy()

import numpy as np
a = np.array([0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1])

ix = np.where(a)[0]
a[ix[:-1]] = np.diff(ix)

print(a[:-1]) # --> array([0, 0, 0, 1, 3, 0, 0, 6, 0, 0, 0, 0, 0, 4, 0, 0, 0])

Upvotes: 1

Frank Yellin
Frank Yellin

Reputation: 11297

Probably not the best answer.

np.where will give you the locations of the non-zero indices in increasing order. By iterating through the result, you know the location of each 1 and the location of the following 1, and can build the result array yourself easily. If the 1s are sparse, this is probably pretty efficient.

Let me see if I can think of something more numpy-ish.

== UPDATE ==

Ah, just came to me

# Find the ones in the array
temp = np.where(x)[0]
# find the difference between adjacent elements
deltas = temp[1:] - temp[:-1]
# Build the result based on these
result = np.zeros_like(x)
result[temp[:-1]] = deltas

Upvotes: 1

Related Questions