Question

假设我有一个2-D NumPy数组，如下所示：

arr = numpy.array([[1,0], [1, 4.6], [2, 10.1], [2, 0], [2, 3.53]])
arr
Out[39]: 
array([[  1.  ,   0.  ],
       [  1.  ,   4.6 ],
       [  2.  ,  10.1 ],
       [  2.  ,   0.  ],
       [  2.  ,   3.53]])

根据第一列中的值对第二列中的值进行分组的最快方法是什么，并从中创建一个dict（所需的输出在下面）

{1: [0, 4.6], 2: [10.1, 0, 3.53]}

目前我使用循环，因为我拥有的实际数组超过100万行，并且第一列有超过5000个唯一值，所以它非常慢。我更喜欢不来使用pandas。

Answer 1

这是一种方法 -

def create_dict(arr):
    a = arr[arr[:,0].argsort()] # sort by col-0 if not already sorted
    s0 = np.r_[0,np.flatnonzero(a[1:,0] > a[:-1,0])+1,a.shape[0]]
    ids = a[s0[:-1],0]
    return {ids[i]:a[s0[i]:s0[i+1],1].tolist() for i in range(len(s0)-1)}

示例运行 -

In [64]: arr
Out[64]: 
array([[  2.  ,   0.  ],
       [  1.  ,   4.6 ],
       [  2.  ,  10.1 ],
       [  4.  ,   0.5 ],
       [  1.  ,   0.  ],
       [  4.  ,   0.23],
       [  2.  ,   3.53]])

In [65]: create_dict(arr)
Out[65]: {1.0: [4.6, 0.0], 2.0: [0.0, 10.1, 3.53], 4.0: [0.5, 0.23]}

运行时测试

其他方法 -

# @Moinuddin Quadri's soln
def defaultdict_based(arr):
    my_list  = arr.tolist()
    my_dict = defaultdict(list)
    for key, value in my_list:
        my_dict[key].append(value)
    return my_dict

# @Psidom's soln
def numpy_split_based(arr):
    sort_arr = arr[arr[:, 0].argsort(), :]
    split_arr = np.split(sort_arr, np.where(np.diff(sort_arr[:,0]))[0] + 1) 
    return {s[0,0]: s[:,1].tolist() for s in split_arr}

计时 -

# Create sample random array with the first col having 1000000 elems
# with 5000 unique ones as stated in the question
In [102]: arr = np.random.randint(0,5000,(1000000,2))

In [103]: %timeit defaultdict_based(arr)
     ...: %timeit numpy_split_based(arr)
     ...: %timeit create_dict(arr)
     ...: 
1 loops, best of 3: 634 ms per loop
1 loops, best of 3: 270 ms per loop
1 loops, best of 3: 260 ms per loop

方法的瓶颈：

似乎与基于defaultdict的方法一样，list与.tolist()的转换证明很重（大约占总运行时间的50％） -

In [104]: %timeit arr.tolist()
1 loops, best of 3: 372 ms per loop

对于其他两种方法，开始时的排序（如果需要）以及最后的分割/循环理解是耗时的部分。排序步骤具有运行时（约占总运行时间的50％） -

In [106]: %timeit arr[arr[:,0].argsort()]
10 loops, best of 3: 140 ms per loop

Answer 2

您可以使用collections.defaultdict在没有numpy的情况下执行此操作。事实上，基于您提供的示例，您甚至不需要numpy数组。 Python list足以满足您的要求。以下是示例：

from collections import defaultdict
my_list = [[1,0], [1, 4.6], [2, 10.1], [2, 0], [2, 3.53]]

my_dict = defaultdict(list)
for key, value in my_list:
    my_dict[key].append(value)

    # if you want the values as float in the dict, use:
    #     my_dict[float(key)].append(float(value))

my_dict保留的最终内容为：

{1: [0, 4.6], 2: [10.1, 0, 3.53]}

Answer 3

您可以使用np.split：

# sort array by the first column if it isn't
sort_arr = arr[arr[:, 0].argsort(), :]

# split the array and construct the dictionary
split_arr = np.split(sort_arr, np.where(np.diff(sort_arr[:,0]))[0] + 1)

{s[0,0]: s[:,1].tolist() for s in split_arr}
# {1.0: [0.0, 4.6], 2.0: [10.1, 0.0, 3.53]}

Answer 4

假设您的第一列按排序顺序，这将有效。

In [165]: d = {}

In [166]: uniq, idx, idxinv, counts = np.unique(arr[:, 0], return_index=True, return_inverse=True, return_counts=True)

In [167]: [d.update({arr[:, 0][el]: arr[:, 1][range(ix, counts[ix])]}) for ix, el in enumerate(idx)]
Out[167]: [None, None]

In [168]: d
Out[168]: {1.0: array([ 0. ,  4.6]), 2.0: array([  4.6,  10.1])}

NumPy：通过将列中的值按其他列值分组来创建dict（）

4 个答案: