CODEWITHSUNDEEP
pythonarraysnumpy
JD80121
JD80121

Reputation: 619

Numpy: vectorized operations to create a 3D array

I am learning Python and would like to find an efficient way to solve this problem using Numpy.

I currently have a 4x8 array containing random integers:

import numpy as np

n = 3
k = np.random.randint(n, size = (4,8))

Each number represents a color defined by its RGB value in a nx3 array:

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

I would like to use these numbers to create a new 4x8x3 array where the first two dimensions represent pixels locations, and the third dimension the color of each pixel. This could be thought of as number painting. For example, if k[3,4] = 2, then myArray[3,4,:] = [255 100 0].

I am getting familiar with Numpy tools, but I am unsure of what I should be looking for exactly. Since the array k will eventually be much larger (I'm thinking ~640x480) and contain more than n = 3 non-random colors, I would like to use vectorized operations in order to speed up the process (and learn a bit more about them). Is this the most efficient way to do it?

Upvotes: 2

Views: 155

Answers (1)

DSM
DSM

Reputation: 353604

IIUC, all you need to do is index into colors with k:

>>> k = np.random.randint(n, size = (2,4))
>>> out = colors[k]
>>> out
array([[[  0, 100, 255],
        [255, 100,   0],
        [255, 100,   0],
        [255, 100,   0]],

       [[  0, 100, 255],
        [  0, 100, 255],
        [255, 100,   0],
        [255, 100,   0]]])
>>> out.shape
(2, 4, 3)
>>> all((out[i]==colors[c]).all() for i,c in np.ndenumerate(k))
True

Upvotes: 1

Related Questions