EN VI

Arrays - Applying a specified rule to build an array?

How to Arrays - Applying a specified rule to build an array

Consider array a whose columns hold random values taken from 1, 2, 3:

a = np.array([[2, 3, 1, 3],
              [3, 2, 1, 3],
              [1, 1, 1, 2],
              [1, 3, 2, 3],
              [3, 3, 1, 3],
              [2, 1, 3, 2]])

Now, consider array b whose first 2 columns hold the 9 possible pairs of values taken from 1, 2, 3 (the order of the pair elements is important). The 3rd column of b associates a non-negative integer with each pairing.

b = np.array([[1, 1, 6],
              [1, 2, 0],
              [1, 3, 9],
              [2, 1, 6],
              [2, 2, 0],
              [2, 3, 4],
              [3, 1, 1],
              [3, 2, 0],
              [3, 3, 8]])

I need help with code that produces array c where vertically adjacent elements in a are replaced with matching values from the 3rd column of b. For example, the first column of 'a' moves down from 2 to 3 to 1 to 1 to 3 to 2. So, the first column of c would hold values 4, 1, 6, 9, 0. The same idea applies to every column of a. We see that pair order is important (moving from 3 to 1 produces value 1, while moving from 1 to 3 produces value 9.

The output of this small example would be:

c = np.array([[4, 0, 6, 8],
              [1, 6, 6, 0],
              [6, 9, 0, 4],
              [9, 8, 6, 8],
              [0, 1, 9, 0]])

Because this code will be executed a vast number of times, I'm hoping there is a speedy vectorized solution. Thanks.

Solution:

Since b contains all the pairs you can efficiently reshape it to square form, indexed by its row/col number, then form the pairs of indices with sliding_window_view and index the square intermediate:

from numpy.lib.stride_tricks import sliding_window_view as swv

s = np.full((b[:, 0].max()+1, b[:, 1].max()+1), -1)
s[b[:, 0], b[:, 1]] = b[:, 2]

v = swv(a, 2, axis=0)
out = s[v[..., 0], v[..., 1]]

Variant with 0 based indexing (produces a slightly more packed intermediate):

s = np.full((b[:, 0].max(), b[:, 1].max()), -1)
s[b[:, 0]-1, b[:, 1]-1] = b[:, 2]

v = swv(a, 2, axis=0)-1
out = s[v[..., 0], v[..., 1]]

Output:

array([[4, 0, 6, 8],
       [1, 6, 6, 0],
       [6, 9, 0, 4],
       [9, 8, 6, 8],
       [0, 1, 9, 0]])

Intermediate s:

array([[-1, -1, -1, -1],
       [-1,  6,  0,  9],
       [-1,  6,  0,  4],
       [-1,  1,  0,  8]])

# variant
array([[6, 0, 9],
       [6, 0, 4],
       [1, 0, 8]])

If you had arbitrary values making it difficult to generate a dense square intermediate, you could use pandas's merge:

from numpy.lib.stride_tricks import sliding_window_view as swv
import pandas as pd

out = (pd.DataFrame(swv(a, 2, axis=0).reshape(-1, 2))
         .merge(pd.DataFrame(b), how='left')
         [2].to_numpy()
         .reshape(-1, a.shape[1])
       )

Output:

array([[4, 0, 6, 8],
       [1, 6, 6, 0],
       [6, 9, 0, 4],
       [9, 8, 6, 8],
       [0, 1, 9, 0]])