0

Hi I have a numpy array like

a = np.array([1,1,1,3,3,6,6,6,6,6,6])

I want to convert it to a continuous array like below

b = np.array([0,0,0,1,1,2,2,2,2,2,2])

I have a code for this using for loop

def fun1(a):
    b = a.copy()
    for i in range(1,a.shape[0]):
        if a[i] != a[i-1]:
            b[i] = b[i-1]+1
        else:
            b[i] = b[i-1]
    b = b - b.min()
    return b

Is there a way to vectorize this using numpy? I can use numba to make it faster but I was wondering if there is way to do it with just numpy.

Scriptable
  • 19,402
  • 5
  • 56
  • 72
deepAgrawal
  • 673
  • 1
  • 7
  • 25

1 Answers1

3

You can do the following with np.unique, using the return_inverse argument. According to the docs:

return_inverse : bool, optional

If True, also return the indices of the unique array (for the specified axis, if provided) that can be used to reconstruct ar.

a = np.array([1,1,1,3,3,6,6,6,6,6,6])
_, b = np.unique(a, return_inverse=True)

>>> b
array([0, 0, 0, 1, 1, 2, 2, 2, 2, 2, 2])

Edit If the list starts out unsorted, and you want such a continuous array as displayed in your output, you can do the same as above, but passing a sorted array to np.unique:

a = np.array([1,3,1,3,3,6,6,3,6,6,6])
_, b = np.unique(sorted(a), return_inverse=True)
>>> b
array([0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 2])
Community
  • 1
  • 1
sacuL
  • 49,704
  • 8
  • 81
  • 106
  • 1
    Whoever down voted, mind to explain? This seems exactly what was asked. – rafaelc May 13 '18 at 03:18
  • 3
    Not the downvoter, but unless the user clarifies that the list will always be sorted, this is not a general solution. Try it for `np.array([3,1,2])` It will produce `[2 0 1]`. The question that this is a dupe of explains this a bit better, and has a general alternative. – user3483203 May 13 '18 at 03:22
  • Thanks @chrisz. I tried to edit it to take that into account. Might still not be the best method, but I think it's a valid option. – sacuL May 13 '18 at 03:32