Skip to content
Advertisement

Replacing chunks of numpy array on condition

Let’s say I have the following numpy array, of 1’s and 0’s exclusively:

import numpy as np

example = np.array([0,1,1,0,1,0,0,1,1], dtype=np.uint8)

I want to group all elements into chunks of 3, and replace the chunks by a single value, based on a condition. Let’s say I want [0,1,1] to become 5, and [0,1,0] to become 10. Thus the desired output would be:

[5,10,5]

All possible combinations of 1’s and 0’s in a chunk have a corresponding unique value that should replace the chunk. What’s the fastest way to do this?

Advertisement

Answer

As the other answers indicated, you can start by reshaping your array (actually, you should probably generate it with the correct shape to begin with, but that’s another issue):

example = np.array([0, 1, 1, 0, 1, 0, 0, 1, 1], dtype=np.uint8)
data = example.reshape(-1, 3)

Now running a custom python function over the array is going to be slow, but luckily numpy has your back here. You can use np.packbits to transform each row into a number directly:

data = np.packbits(data, axis=1, bitorder='little').ravel() # [6, 2, 6]

If you wanted 101 to map to 5 and 110 to map to 6, your job is done. Otherwise, you will need to come up with a mapping. Since you have three bits, you only need 8 numbers in the mapping array:

mapping = np.array([7, 4, 3, 8, 124, 1, 5, 0])

You can use data as an index directly into mapping. The output will have the type of mapping but the shape of data:

result = mapping[data]  # [5, 3, 5]

You can do this in one line:

mapping[np.packbits(example.reshape(-1, 3), axis=1, bitorder='little').ravel()]
User contributions licensed under: CC BY-SA
9 People found this is helpful
Advertisement