Snake traversal of 2D NumPy array

Question

I have the following 2D array:

In [173]: arr
Out[173]: 
array([[ 1,  2,  3,  4],   # -> -> -> ->
       [ 5,  6,  7,  8],   # <- <- <- <-
       [ 9, 10, 11, 12],   # -> -> -> ->
       [13, 14, 15, 16],   # <- <- <- <-
       [17, 18, 19, 20]])  # -> -> -> ->

And I'd like to traverse the array in a snake-like pattern starting from top left element and ending up in bottom right element.

As of now, I have this uninteresting way of solving:

In [187]: np.hstack((arr[0], arr[1][::-1], arr[2], arr[3][::-1], arr[4]))
Out[187]: 
array([ 1,  2,  3,  4,  8,  7,  6,  5,  9, 10, 11, 12, 16, 15, 14, 13, 17,
       18, 19, 20])

How can we do it with minimal effort, without looping and not too much hard-coding?

Answer 1

One way would be starting off with a copy of the input and then replacing the second row onwards with column flipped version of the corresponding rows off the input and doing this for all the even rows using step-sized slicing. Finally, a ravel() is needed at the end for the desired flattened version.

Hence, the implementation would look something like this -

out = arr.copy()
out[1::2] = arr[1::2,::-1]
out = out.ravel()

---

Another compact way would be with np.where to do the chosing between col-flipped and non-flipped versions and hence achieve our desired output -

np.where(np.arange(len(arr))[:,None]%2,arr[:,::-1],arr).ravel()

Explanation with the given sample -

# Array to be used for the chosing. 1s would be True ones and 0s are False
In [72]: np.arange(len(arr))[:,None]%2
Out[72]: 
array([[0],
       [1],
       [0],
       [1],
       [0]])

# Use np.where to choose. So, arr[:,::-1] must be the first data, as
# that's one to be put on even rows and arr would be the second one.
In [73]: np.where(np.arange(len(arr))[:,None]%2,arr[:,::-1],arr)
Out[73]: 
array([[ 1,  2,  3,  4],
       [ 8,  7,  6,  5],
       [ 9, 10, 11, 12],
       [16, 15, 14, 13],
       [17, 18, 19, 20]])

# Finally flatten
In [74]: np.where(np.arange(len(arr))[:,None]%2,arr[:,::-1],arr).ravel()
Out[74]: 
array([ 1,  2,  3,  4,  8,  7,  6,  5,  9, 10, 11, 12, 16, 15, 14, 13, 17,
       18, 19, 20])