Best way to find indexes of first occurences of integers in each row of numpy array?

user_14693068 0
Viewed 75

If I have an array such as:

a = np.array([[1, 1, 2, 2, 1, 3, 4],
              [8, 7, 7, 7, 4, 8, 8]])

what would be the best way to get as output:

array([[0, 2, 5, 6], [0, 1, 4]])

or

array([[0, 2, 5, 6], [4, 1, 0]])

These are the indices of the first occurence of each integer in each row. The order of the indices is not important.

Currently I am using:

res = []
for row in a:
  unique, unique_indexes = np.unique(a, return_index=True)
  res.append(unique_indexes)

But I wonder if there is a (num)pythonic way to avoid the for loop.

numpy python
2 Answers

You can transform the array in such a way that you process the entire thing in one batch. Let's start with an example very similar to the one in your question:

a = np.array([[1, 1, 2, 2, 1, 3, 4], [8, 7, 7, 7, 5, 8, 8]])

Now get the indices:

_, ix = np.unique(a, return_index=True)
# ix = array([ 0,  2,  5,  6, 11,  8,  7])

Notice that the indices of the first elements are correct. The following elements are offset by the size of a. In general, the offset is

offset = ix // a.shape[-1]
# offset = array([0, 0, 0, 0, 1, 1, 1])
ix %= a.shape[-1]
# ix = array([0, 2, 5, 6, 4, 1, 0])

You can call np.split on the new ix at every location where offset changes value:

ix = np.split(ix, np.flatnonzero(np.diff(offset)) + 1)

So why is this example valid, but the one in the question is not? The key is that np.unique uses a sort-based approach (which makes it run in O(n log n) rather than the O(n) of collections.Counter). That means that for the order of the indices to be correct, each row must be unique from and larger than the previous row. Notice that in your example, 4 appears in both rows. You can ensure this with a simple check of the max and min values in each row:

mn = a.min(axis=1)
mx = a.max(axis=1)
diff = np.r_[0, (mx - mn + 1)[:-1].cumsum(0)] - mn
# diff = array([-4, -4])
b = a + diff[:, None]
# b = array([[0, 0, 1, 1, 0, 2, 3],
             [7, 6, 6, 6, 4, 7, 7]])

Notice that you have to offset the cumulative sum by one to get the right index. If you deal with large integers and/or very large arrays, you will have to be more careful about making diff to avoid overflow.

Now you can use b in place of a in the call to np.unique.

TL;DR

Here is a general no-loop approach, applicable to any axis, not just the last:

def global_unq(a, axis=-1):
    n = a.shape[axis]
    a = np.moveaxis(np.asanyarray(a), axis, -1).reshape(-1, n)
    mn = a.min(-1)
    mx = a.max(-1)
    diff = np.r_[0, (mx - mn + 1)[:-1].cumsum(0)] - mn
    _, ix = np.unique(a + diff[:, None], return_index=True)
    return np.split(ix % n, np.flatnonzero(np.diff(ix // n)) + 1)
user_29887300

You could put it into a list comprehension, but your loop is fairly clean already:

[list(np.unique(e, return_index=True)[-1]) for e in a]
# [[0, 2, 5, 6], [4, 1, 0]]
user_119717850
Related