python - Indexing a list with an unique index

Question

I have a list say l = [10,10,20,15,10,20]. I want to assign each unique value a certain "index" to get [1,1,2,3,1,2].

This is my code:

a = list(set(l))
res = [a.index(x) for x in l]

Which turns out to be very slow.

l has 1M elements, and 100K unique elements. I have also tried map with lambda and sorting, which did not help. What is the ideal way to do this?

Answer 1

You can do this in O(N) time using a defaultdict and a list comprehension:

>>> from itertools import count
>>> from collections import defaultdict
>>> lst = [10, 10, 20, 15, 10, 20]
>>> d = defaultdict(count(1).next)
>>> [d[k] for k in lst]
[1, 1, 2, 3, 1, 2]

In Python 3 use __next__ instead of next.

---

If you're wondering how it works?

The default_factory(i.e count(1).next in this case) passed to defaultdict is called only when Python encounters a missing key, so for 10 the value is going to be 1, then for the next ten it is not a missing key anymore hence the previously calculated 1 is used, now 20 is again a missing key and Python will call the default_factory again to get its value and so on.

d at the end will look like this:

>>> d
defaultdict(<method-wrapper 'next' of itertools.count object at 0x1057c83b0>,
            {10: 1, 20: 2, 15: 3})