如何从二维Numpy数组中反向检索最初用于构造它的原始列表

原始场景说明 #

我们有以下由4个列表生成笛卡尔积构造的numpy数组：

import numpy as np

a1 = [1,2,3,4]
a2 = [11,22,33,44]
a3 = [111,222,333,444]
a4 = [1111,2222,3333,4444]

a = []
for x in a1:
    for y in a2:
        for k in a3:
            for l in a4:
                a.append((x, y, k, l))

na = np.array(a)

你目前的实现通过手动reshape数组维度后逐维索引提取原始列表，运行符合预期：

na.shape = (4,4,4,4,4)
a1 = na[:,0,0,0,0]
a2 = na[0,:,0,0,1]
a3 = na[0,0,:,0,2]
a4 = na[0,0,0,:,3]

更简洁的实现方案 #

方案1：按列去重提取（通用性最强） #

因为笛卡尔积生成的二维数组，每一列对应原始列表全量元素的重复值，直接对每一列去重即可得到原始列表，不需要手动调整数组维度：

# 若原始列表本身是升序排列，可直接用unique
a1, a2, a3, a4 = [np.unique(na[:, i]) for i in range(4)]

# 若需要严格保留原始列表的顺序、不受排序影响，使用return_index参数
def get_ordered_unique(arr, col):
    _, idx = np.unique(arr[:, col], return_index=True)
    return arr[:, col][np.sort(idx)]

a1, a2, a3, a4 = [get_ordered_unique(na, i) for i in range(4)]

该方案不需要提前知道原始列表的长度，也不需要修改数组shape，适配任意数量原始列表生成的笛卡尔积数组。

方案2：简化reshape索引写法 #

如果坚持用reshape的方式，可通过语法简化多维度的0索引，减少冗余代码：

na = na.reshape(4,4,4,4,4)
a1 = na[:, *[0]*3, 0]
a2 = na[0, :, *[0]*2, 1]
a3 = na[*[0]*2, :, 0, 2]
a4 = na[*[0]*3, :, 3]

内容的提问来源于stack exchange，提问作者Damir Devetak