我怎样才能快速花式重新排序扁平化的 "jagged" numpy 数组

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

所以我在一个平面数组中有大量数据,这些数据被分组为不规则大小的块。这些块的大小在另一个数组中给出。我需要做的是根据第三个索引数组重新排列块(想想花哨的索引)

这些块的长度总是 >= 3,通常是 4,但技术上没有限制,因此填充到最大长度和掩码是不可行的。此外,由于技术原因,我只能访问 numpy,所以没有像 scipy 或 pandas.

为了便于阅读,本示例中的数据很容易分组。在真实数据中,数字可以是任何东西,不遵循这种模式。

[编辑] 更新了较少混淆的数据

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])

这种情况下的预期输出是

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

由于真实数据可能长达数百万,我希望有某种 numpy 魔法可以在没有 python 循环的情况下完成此操作。

如果您使用 np.cumsum 构建索引,则可以使用 np.split 在对应于 chunkSizes 的 data 数组中创建视图。然后,您可以使用花式索引根据 newOrder 索引对视图重新排序。这应该是相当有效的,因为当您在重新排序的视图上调用 np.concatenate 时,数据只会复制到新数组:

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])

方法 #1

这是基于创建常规数组和掩码的矢量化 -

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]]

给定样本的输出 -

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])

方法 #2

另一个基于 cumsum 的矢量化,并且对于那些 非常破烂 块大小的占用空间更小 -

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