CODEWITHSUNDEEP
pythonnumpy
Dzung Nguyen
Dzung Nguyen

Reputation: 3942

Indexing in numpy array

Given a lookup table:

colors = [  [0,0,0],\
            [0, 255, 0],\
            [0, 0, 255],\
            [255, 0, 0]]

And an index numpy matrix input 2x2:

a = np.array([[0,1],[1,1]])

How can I map a to a 2x2x3 matrix b, where b[i][j] = colors[a[i][j]]? I want to avoid using a for loop here.

Upvotes: 1

Views: 109

Answers (1)

Bi Rico
Bi Rico

Reputation: 25833

Have you tried:

colors[a]

Here's a full example:

import numpy as np

colors = np.array([[0,0,0],
                   [0, 255, 0],
                   [0, 0, 255],
                   [255, 0, 0]
                  ])
a = np.array([[0, 1], [1, 1]])

new = colors[a]
new.shape
# (2, 2, 3)
new
# array([[[  0,   0,   0],
#         [  0, 255,   0]],
#
#        [[  0, 255,   0],
#         [  0, 255,   0]]])

Upvotes: 2

Related Questions