Map NumPy integer list values to lists according to lambda function

Viewed 71

How can I map each value of an integer NumPy array to a list according to a lambda function?

I have the following code:

x = [0, 1, 2, 0, 0, 2, 1, 1]
colours = [[0, 0, 0], [255, 0, 0], [0, 255, 0], [0, 0, 255]]

coloured = lambda v: colours[int(v)]
vcoloured = np.vectorize(coloured)

x_color = vcoloured(x) # throws ValueError: setting an array element with a sequence.

# would like to get x_color = [[0, 0, 0], [255, 0, 0], [0, 255, 0], [0, 0, 0], [0, 0, 0], [0, 255, 0], [255, 0, 0], [255, 0, 0]]
2 Answers

Don't bother with lists at all. If you have numpy arrays, it's much faster to use the power of numpy:

>>> x = np.array([0, 1, 2, 0, 0, 2, 1, 1])
>>> colors = np.array([[0, 0, 0], [255, 0, 0], [0, 255, 0], [0, 0, 255]])
>>> x_color = colors[x]
>>> x_color
array([[  0,   0,   0],
       [255,   0,   0],
       [  0, 255,   0],
       [  0,   0,   0],
       [  0,   0,   0],
       [  0, 255,   0],
       [255,   0,   0],
       [255,   0,   0]])

Yeah, I went ahead and Americanized your spelling, but functionally, the array x_colored will behave just like a list for all practical purposes. If you absolutely need a list, just call tolist on the result:

x_color = colors[x].tolist()

Your use of vectorized has problems because it has to guess as to the returned array dtype. Your function returns a list, not a number. So we need to specify the return dtype - as object:

In [25]: x = [0, 1, 2, 0, 0, 2, 1, 1]
    ...: colours = [[0, 0, 0], [255, 0, 0], [0, 255, 0], [0, 0, 255]]
    ...: 
    ...: coloured = lambda v: colours[int(v)]
    ...: vcoloured = np.vectorize(coloured, otypes=[object])
In [26]: vcoloured(x)
Out[26]: 
array([list([0, 0, 0]), list([255, 0, 0]), list([0, 255, 0]),
       list([0, 0, 0]), list([0, 0, 0]), list([0, 255, 0]),
       list([255, 0, 0]), list([255, 0, 0])], dtype=object)

That could be turned into a 2d array with:

In [27]: np.stack(_)
Out[27]: 
array([[  0,   0,   0],
       [255,   0,   0],
       [  0, 255,   0],
       [  0,   0,   0],
       [  0,   0,   0],
       [  0, 255,   0],
       [255,   0,   0],
       [255,   0,   0]])

But vectorize is not a performance tool (read its docs). It's faster to use a list comprehension:

In [28]: np.array([colours[v] for v in x])
Out[28]: 
array([[  0,   0,   0],
       [255,   0,   0],
       [  0, 255,   0],
       [  0,   0,   0],
       [  0,   0,   0],
       [  0, 255,   0],
       [255,   0,   0],
       [255,   0,   0]])

The direct array indexing of @Mad's answer is better.

Related