Unzigzag an array into a matrix-CodePudding

I have a 1D array which is in fact a 2D matrix sampled this way:

↓-------<------S  <- start
>------->------↓
↓-------<------<
>------->------E  <- end

For example

B =    1 2  3  4
       5 6  7  8
       9 10 11 12

is coded as A = [4, 3, 2, 1, 5, 6, 7, 8, 12, 11, 10, 9].
The number of rows can be odd or even.

The following code works but is inefficient due to the loop (and no vectorization). How to do a more efficient Numpy "unzigzaging"?

import numpy as np

def unzigzag(z, numcols):
    numrows = len(z) // numcols
    a = np.zeros((numrows, numcols))
    col = numcols - 1
    row = 0
    sign = -1
    for i in range(numrows*numcols):
        if col == -1:
            col = 0
            sign = 1
            row  = 1
        if col == numcols:
            col = numcols - 1
            sign = -1
            row  = 1
        a[row, col] = z[i]
        col  = sign
    return a    

A = [4, 3, 2, 1, 5, 6, 7, 8, 12, 11, 10, 9]
B = unzigzag(A, 4)
#[[ 1.  2.  3.  4.]
# [ 5.  6.  7.  8.]
# [ 9. 10. 11. 12.]]

If possible, it would be useful to have it working even if more dimensions than 1D:

if A has shape (12, ), unzigzag(A, numcols=4) has shape (3, 4)
if A has shape (12, 100), unzigzag(A, numcols=4) has shape (3, 4, 100)
if A has shape (n, i, j, k, ...), unzigzag(A, numcols) has shape (n/numcols, numcols, i, j, k, ...)

Edit: example for the n-dimensional case:

import numpy as np

def unzigzag3(z, numcols):
        numrows = z.shape[0] // numcols
        new_shape = (numrows, numcols)   z.shape[1:]
        a = np.zeros(new_shape)
        col = numcols - 1
        row = 0
        sign = -1
        for i in range(numrows*numcols):
            if col == -1:
                col = 0
                sign = 1
                row  = 1
            if col == numcols:
                col = numcols - 1
                sign = -1
                row  = 1
            a[row, col, :] = z[i, :]
            col  = sign
        return a    

A = np.array([[4,4], [3, 3], [2, 2], [1, 1], [5, 5], [6, 6], [7, 7], [8, 8], [12, 12], [11, 11], [10, 10], [9, 9]])
print(A.shape)  # (12, 2)
B = unzigzag3(A, 4)
print(B) 
print(B.shape) # (3, 4, 2)

CodePudding user response：

I would reshape, then updated every other row with the horizontally flipped version of that row.

import numpy as np

a = np.array([4, 3, 2, 1, 5, 6, 7, 8, 12, 11, 10, 9])
numcols = 4

a = a.reshape(-1,numcols)
a[::2] = np.flip(a, axis=1)[::2]

print(a)

Output

[[ 1  2  3  4]
 [ 5  6  7  8]
 [ 9 10 11 12]]

CodePudding user response：

Based on @Chris' answer (full credit to him), here seems to be a working version for both 1D and n-dimensional input:

import numpy as np

def unzigzag(z, numcols):
    new_shape = (z.shape[0] // numcols, numcols)   z.shape[1:]
    a = z.copy().reshape(new_shape)
    a[::2] = np.flip(a, axis=1)[::2]
    return a

unzigzag(np.array([4, 3, 2, 1, 5, 6, 7, 8, 12, 11, 10, 9]), numcols=4)

# [[ 1  2  3  4]
#  [ 5  6  7  8]
#  [ 9 10 11 12]]
# shape: (3, 4)

unzigzag(np.array([[4,4], [3, 3], [2, 2], [1, 1], [5, 5], [6, 6], [7, 7], [8, 8], [12, 12], [11, 11], [10, 10], [9, 9]]), numcols=4)

# [[[ 1  1]
#   [ 2  2]
#   [ 3  3]
#   [ 4  4]]

#  [[ 5  5]
#   [ 6  6]
#   [ 7  7]
#   [ 8  8]]

#  [[ 9  9]
#   [10 10]
#   [11 11]
#   [12 12]]]
# shape: (3, 4, 2)