如何根据列表对2D numpy对象数组进行排序

Question

我有一个二维的numpy对象数组：

aa = np.array([["aaa","05","1","a"],
               ["ccc","30","2","v"],
               ["ddd","50","2","v"],
               ["bbb","10","1","v"]])

以及以下列表：

sample_ids = ["aaa", "bbb", "ccc", "ddd"]

我想根据列表对numpy数组进行排序，以便获得以下结果：

[["aaa","05","1","a"],
 ["bbb","10","1","v"],
 ["ccc","30","2","v"],
 ["ddd","50","2","v"]]
 

编辑：

如果（在sample_ids中）存在不在数组中的键，则结果数组将不包括这些缺失的键（即不添加空行）。 因此，如果我们有以下内容：

sample_ids = ["aaa", "bbb", "ccc", "ddd", "eee"]

最终数组仍将保持不变。 此外，如果该数组包含一个键（即行键）在键中不存在，那么该行也将从最终数组中省略。 编辑2： 从尼克的答案开始，我想出了这个处理缺失键的方法。

sample_ids2 = ["aaa", "bbb", "eee", "ccc", "ddd"]

idxs = []
for i,v in enumerate(sample_ids2):
    if str(list(aa.T[0])).find(v) != -1:
        k = list(aa.T[0]).index(v)
        idxs.append(k)
    else:
        print(v + " was not found!!!")

print(aa[idxs])

输出：

[['aaa' '05' '1' 'a']
 ['bbb' '10' '1' 'v']
 ['ccc' '30' '2' 'v']
 ['ddd' '50' '2' 'v']]

- julio514

---

我有两个问题：首先，您的数据是一个字符串表格，您是否想使用更擅长处理此类数据的numpy而不是pandas来解决它？其次，我们需要考虑重复/缺失键吗？（例如 sample_ids = ["aaa", "aaa", "ddd"]） - ken

非常感谢大家的输入。 我现在会继续使用numpy。实际上，数组中可能缺少键（sample_ids）。在这种情况下，这些键需要从结果排序后的数组中省略。更多详细信息请参见编辑。 - julio514

Answer 1

以下是一些可能的解决方案。使用numpy库：

subs = list(aa.T[0])
idxs = [subs.index(i) for i in sample_ids if i in subs]
res = aa[idxs]
# array([['aaa', '05', '1', 'a'],
#        ['bbb', '10', '1', 'v'],
#        ['ccc', '30', '2', 'v'],
#        ['ddd', '50', '2', 'v']], dtype='<U3')

使用pandas：

res = np.array(pd.DataFrame(aa).set_index(0).reindex(sample_ids).dropna().reset_index())
# array([['aaa', '05', '1', 'a'],
#        ['bbb', '10', '1', 'v'],
#        ['ccc', '30', '2', 'v'],
#        ['ddd', '50', '2', 'v']], dtype=object)

对于这两种情况，如果sample_ids = ["aaa", "bbb", "ccc", "ddd", "eee"]，输出将是相同的。

如果sample_ids = ["ddd", "aaa", "bbb"]，输出将是：

array([['ddd', '50', '2', 'v'],
       ['aaa', '05', '1', 'a'],
       ['bbb', '10', '1', 'v']])

- Nick

---

@julio514 我已经更新了我的解决方案，使其在aa或sample_ids中存在另一个中没有的值时也能正常工作。 - Nick

Answer 2

受@Nick的第一种方法启发：

# first build a dictionary of value: position
key = {k: i for i, k in enumerate(sample_ids)}
# {'aaa': 0, 'bbb': 1, 'ccc': 2, 'ddd': 3}

# then sort based on this key
out = aa[np.argsort(np.vectorize(key.get)(aa[:, 0]))]

如果你想要处理缺失值（默认键为-1，如果你喜欢最后使用np.inf）：

out = aa[np.argsort(np.vectorize(lambda x: key.get(x, -1))(aa[:, 0]))]

输出：

array([['aaa', '05', '1', 'a'],
       ['bbb', '10', '1', 'v'],
       ['ccc', '30', '2', 'v'],
       ['ddd', '50', '2', 'v']], dtype='<U3')

- mozway

Answer 3

一个可能的解决方案，使用 numpy.broadcasting：

aa[np.argmax(aa[:,0] == np.array(sample_ids)[:, None], axis=1), :]

另一种可能的解决方案，使用列表推导式：

np.vstack([aa[aa[:,0] == x, :] for x in sample_ids])

输出：

array([['aaa', '05', '1', 'a'],
       ['bbb', '10', '1', 'v'],
       ['ccc', '30', '2', 'v'],
       ['ddd', '50', '2', 'v']], dtype='<U3')

- PaulS