python - 对非常大的 numpy 数组按 ID 进行分组的最快方法

Question

我正在尝试找到对具有相似 ID 的“行”进行分组的最佳方法。

我最好的猜测: np.array([test[test[:,0] == ID] for ID in List_IDs])

结果:数组的数组的数组

[ array([['ID_1', 'col1','col2',...,'coln'],
         ['ID_1', 'col1','col2',...,'coln'],...,
         ['ID_1', 'col1','col2',...,'coln']],dtype='|S32')
array([['ID_2', 'col1','col2',...,'coln'],
         ['ID_2', 'col1','col2',...,'coln'],...,
         ['ID_2', 'col1','col2',...,'coln']],dtype='|S32')
....
array([['ID_k', 'col1','col2',...,'coln'],
         ['ID_k', 'col1','col2',...,'coln'],...,
         ['ID_K', 'col1','col2',...,'coln']],dtype='|S32')]

任何人都可以建议一些更有效的方法吗？

提醒:test 数组很大。 “行”未排序

Answer 1

我假设 List_IDs 是第一列中所有唯一 ID 的列表。有了这个假设，这是一个基于 Numpy 的解决方案 -

# Sort input array test w.r.t. first column that are IDs
test_sorted = test[test[:,0].argsort()]

# Convert the string IDs to numeric IDs
_,numeric_ID = np.unique(test_sorted[:,0],return_inverse=True)

# Get the indices where shifts (IDs change) occur
_,cut_idx = np.unique(numeric_ID,return_index=True)

# Use the indices to split the input array into sub-arrays with common IDs
out = np.split(test_sorted,cut_idx)[1:]

示例运行 -

In [305]: test
Out[305]: 
array([['A', 'A', 'B', 'E', 'A'],
       ['B', 'E', 'A', 'E', 'B'],
       ['C', 'D', 'D', 'A', 'C'],
       ['B', 'D', 'A', 'C', 'A'],
       ['B', 'A', 'E', 'A', 'E'],
       ['C', 'D', 'C', 'E', 'D']], 
      dtype='|S32')

In [306]: test_sorted
Out[306]: 
array([['A', 'A', 'B', 'E', 'A'],
       ['B', 'E', 'A', 'E', 'B'],
       ['B', 'D', 'A', 'C', 'A'],
       ['B', 'A', 'E', 'A', 'E'],
       ['C', 'D', 'D', 'A', 'C'],
       ['C', 'D', 'C', 'E', 'D']], 
      dtype='|S32')

In [307]: out
Out[307]: 
[array([['A', 'A', 'B', 'E', 'A']], 
       dtype='|S32'), array([['B', 'E', 'A', 'E', 'B'],
        ['B', 'D', 'A', 'C', 'A'],
        ['B', 'A', 'E', 'A', 'E']], 
       dtype='|S32'), array([['C', 'D', 'D', 'A', 'C'],
        ['C', 'D', 'C', 'E', 'D']], 
       dtype='|S32')]