How can I quickly fancy-reorder a flattened "jagged" numpy array

user_654602 0
Viewed 196

So I have lots of data in a single, flat array that is grouped into irregularly sized chunks. The sizes of these chunks are given in another array. What I need to do is rearrange the chunks based on a third index array (think fancy indexing)

These chunks are always >= 3 long, usually 4, but technically unbounded, so it's not feasible to pad up to a max length and mask. Also, due to technical reasons I only have access to numpy, so nothing like scipy or pandas.

Just to be easier to read, the data in this example is easily grouped. In the real data, the numbers can be anything and do not follow this pattern.

[EDIT] Updated with less confusing data

data = np.array([1,2,3,4, 11,12,13, 21,22,23,24, 31,32,33,34, 41,42,43, 51,52,53,54])
chunkSizes = np.array([4, 3, 4, 4, 3, 4])
newOrder = np.array([0, 5, 4, 5, 2, 1])

The expected output in this case would be

np.array([1,2,3,4, 51,52,53,54, 41,42,43, 51,52,53,54, 21,22,23,24, 11,12,13])

Since the real data can be millions long, I'm hoping for some kind of numpy magic that can do this without python loops.

jagged-arrays numpy python performance
2 Answers

Approach #1

Here's a vectorized one based on creating a regular array and masking -

def chunk_rearrange(data, chunkSizes, newOrder):
    m = chunkSizes[:,None] > np.arange(chunkSizes.max())
    d1 = np.empty(m.shape, dtype=data.dtype)
    d1[m] = data
    return d1[newOrder][m[newOrder]]

Output for given sample -

In [4]: chunk_rearrange(data, chunkSizes, newOrder)
Out[4]: array([0, 0, 0, 0, 5, 5, 5, 5, 4, 4, 4, 5, 5, 5, 5, 2, 2, 2, 2, 1, 1, 1])

Approach #2

Another vectorized one based on cumsum and with smaller footprint for those very-ragged chunksizes -

def chunk_rearrange_cumsum(data, chunkSizes, newOrder):
    # Setup ID array that will hold specific values at those interval starts,
    # such that a final cumsum would lead us to the indices which when indexed
    # by the input array gives us the re-arranged o/p   
    idar = np.ones(len(data), dtype=int)

    # New chunk lengths
    newlens = chunkSizes[newOrder]

    # Original chunk intervals
    c = np.r_[0,chunkSizes[:-1].cumsum()]

    # Indices from original order that form the interval starts in new arrangement
    d1 = c[newOrder]

    # Starts of chunks in new arrangement where those from d1 are to be assigned
    c2 = np.r_[0,newlens[:-1].cumsum()]

    # Offset required for the starts in new arrangement for final cumsum to work
    diffs = np.diff(d1)+1-np.diff(c2)
    idar[c2[1:]] = diffs
    idar[0] = d1[0]

    # Final cumsum and indexing leads to desired new arrangement
    out = data[idar.cumsum()]
    return out
user_32938810

You can use np.split to create views into your data array corresponding to the chunkSizes, if you build up the indices with np.cumsum. You can then reorder the views according to the newOrder indices using fancy indexing. This should be reasonably efficient since the data is only copied to the new array when you call np.concatenate on the reordered views:

import numpy as np

data = np.array([0,0,0,0, 1,1,1, 2,2,2,2, 3,3,3,3, 4,4,4, 5,5,5,5])
chunkSizes = np.array([4, 3, 4, 4, 3, 4])
newOrder = np.array([0, 5, 4, 5, 2, 1])

cumIndices = np.cumsum(chunkSizes)
splitArray = np.array(np.split(data, cumIndices[:-1]))
targetArray = np.concatenate(splitArray[newOrder])

# >>> targetArray
# array([0, 0, 0, 0, 5, 5, 5, 5, 4, 4, 4, 5, 5, 5, 5, 2, 2, 2, 2, 1, 1, 1])
user_136182390
Related