How to get a reverse mapping in numpy in O(1)?

Viewed 2401

I have a numpy array, whose elements are unique, for example:

b = np.array([5, 4, 6, 8, 1, 2])

(Edit2: b can have large numbers, and float numbers. The above example is there for simplicity)

I am getting numbers, that are elements in b.

I want to find their index in b, meaning I want a reverse mapping, from value to index, in b.

I could do

for number in input:
    ind = np.where(number==b)

which would iterate over the entire array every call to where.

I could also create a dictionary,

d = {}
for i, element in enumerate(list(b)):
    d[element] = i

I could create this dictionary at "preprocessing" time, but still I would be left with a strange looking dictionary, in a mostly numpy code, which seems (to me) not how numpy is meant to be used.

How can I do this reverse mapping in numpy?

usage (O(1) time and memory required):

print("index of 8 is: ", foo(b, 8))

  • Edit1: not a duplicate of this

Using in1d like explained here doesn't solve my problem. Using their example:

b = np.array([1, 2, 3, 10, 4])

I want to be able to find for example 10's index in b, at runtime, in O(1).

Doing a pre-processing move

mapping = np.in1d(b, b).nonzero()[0]

>> [0, 1, 2, 3, 4]

(which could be accomplished using np.arange(len(b)))

doesn't really help, because when 10 comes in as input, It is not possible to tell its index in O(1) time with this method.

4 Answers

It's simpler than you think, by exploiting numpy's advanced indexing.

What we do is make our target array and just assign usign b as an index. We'll assign the indices we want by using arange.

>>> t = np.zeros((np.max(b) + 1,))
>>> t[b] = np.arange(0, b.size)
>>> t
array([0., 4., 5., 0., 1., 0., 2., 0., 3.])

You might use nans or -1 instead of zeros to construct the target to help detect invalid lookups.

Memory usage: this is optimally performant in both space and time as it's handled entirely by numpy.

If you can tolerate collisions, you can implement a poor man's hashtable. Suppose we have currencies, for example:

h = np.int32(b * 100.0) % 101  # Typically some prime number
t = np.zeros((101,))
t[h] = np.arange(0, h.size)

# Retrieving a value v; keep in mind v can be an ndarray itself.
t[np.int32(v * 100.0) % 101]

You can do any other steps to munge the address if you know what your dataset looks like.

This is about the limit of what's useful to do with numpy.

Solution

If you want constant time (ie O(1)), then you'll need to precompute a lookup table of some sort. If you want to make your lookup table using another Numpy array, it'll effectively have to be a sparse array, in which most values are "empty". Here's a workable approach in which empty values are marked as -1:

b = np.array([5, 4, 6, 8, 1, 2])

_b_ix = np.array([-1]*(b.max() + 1))
_b_ix[b] = np.arange(b.size)
# _b_ix: array([-1,  4,  5, -1,  1,  0,  2, -1,  3])

def foo(*val):
    return _b_ix[list(val)]

Test:

print("index of 8 is: %s" % foo(8))
print("index of 0,5,1,8 is: %s" % foo(0,5,1,8))

Output:

index of 8 is: [3]
index of 0,5,1,8 is: [-1  0  4  3]

Caveat

In production code, you should definitely use a dictionary to solve this problem, as other answerers have pointed out. Why? Well, for one thing, say that your array b contains float values, or any non-int value. Then a Numpy-based lookup table won't work at all.

Thus, you should use the above answer only if you have a deep-seated philosophical opposition to using a dictionary (eg a dict ran over your pet cat). Here's a nice way to generate a reverse lookup dict:

ix = {k:v for v,k in enumerate(b.flat)}

You can use dict, zip and numpy.arrange to create your reverse lookup:

import numpy 

b = np.array([5, 4, 6, 8, 1, 2])
d = dict(zip(b, np.arange(0,len(b))))
print(d)

gives:

{5: 0, 4: 1, 6: 2, 8: 3, 1: 4, 2: 5}

If you want to do multiple lookups, you can do these in O(1) after an initial O(n) traversal to create a lookup dictionary.

b = np.array([5, 4, 6, 8, 1, 2])
lookup_dict = {e:i for i,e in enumerate(b)}
def foo(element):
    return lookup_dict[element]

And this works for your test:

>>> print('index of 8 is:', foo(8))
index of 8 is:  3

Note that if there is a possibility that b may have changed since the last foo() call, we must re-create the dictionary.

Related