形状不匹配的NumPy阵列交织

Question

我想沿着特定的轴交错具有不同尺寸的多个numpy数组。特别是，我有一个(_, *dims)沿第一轴变化的形状数组的列表，我想对其进行交织以获得另一个形状数组(_, *dims)。例如，给定输入

a1 = np.array([[11,12], [41,42]])
a2 = np.array([[21,22], [51,52], [71,72], [91,92], [101,102]])
a3 = np.array([[31,32], [61,62], [81,82]])

interweave(a1,a2,a3)

所需的输出将是

np.array([[11,12], [21,22], [31,32], [41,42], [51,52], [61,62], [71,72], [81,82], [91,92], [101,102]]

借助先前的文章（例如Numpy带有交错的串联数组），当数组沿第一个维度匹配时，我就可以正常工作了：

import numpy as np

def interweave(*arrays, stack_axis=0, weave_axis=1):
    final_shape = list(arrays[0].shape)
    final_shape[stack_axis] = -1

    # stack up arrays along the "weave axis", then reshape back to desired shape
    return np.concatenate(arrays, axis=weave_axis).reshape(final_shape)

不幸的是，如果输入形状沿第一维不匹配，则上面的方法会引发异常，因为我们必须沿着与不匹配的轴不同的轴进行连接。确实，我在这里看不到任何有效使用串联的方法，因为沿着不匹配的轴进行串联将破坏我们生成所需输出所需的信息。

我的另一个想法是用空条目填充输入数组，直到它们的形状沿第一个维度匹配，然后在一天结束时删除空条目。尽管这会起作用，但我不确定如何最好地实现它，而且似乎一开始就没有必要这样做。

Answer 1

这是一种主要NumPy基于的方法，还使用zip_longest填充值来交错数组：

def interleave(*a):
    # zip_longest filling values with as many NaNs as
    # values in second axis
    l = *zip_longest(*a, fillvalue=[np.nan]*a[0].shape[1]),
    # build a 2d array from the list
    out = np.concatenate(l)
    # return non-NaN values
    return out[~np.isnan(out[:,0])]

a1 = np.array([[11,12], [41,42]])
a2 = np.array([[21,22], [51,52], [71,72], [91,92], [101,102]])
a3 = np.array([[31,32], [61,62], [81,82]])

interleave(a1,a2,a3)

array([[ 11.,  12.],
       [ 21.,  22.],
       [ 31.,  32.],
       [ 41.,  42.],
       [ 51.,  52.],
       [ 61.,  62.],
       [ 71.,  72.],
       [ 81.,  82.],
       [ 91.,  92.],
       [101., 102.]])