时间戳

首页 > 解决方案 > numpy 根据唯一列创建子数组

python - numpy 根据唯一列创建子数组

问题描述

我有一个示例数组，看起来像array = np.array([[1,1,0,1], [0,1,0,0], [1,1,1,0], [0,0,1,2], [0,1,3,2], [1,1,0,1], [0,1,0,0]])......

array([[1, 1, 0, 1],
       [0, 1, 0, 0],
       [1, 1, 1, 0],
       [0, 0, 1, 2],
       [0, 1, 3, 2],
       [1, 1, 0, 1],
       [0, 1, 0, 0]])

考虑到这一点，我想根据前两列将此数组重新格式化为子数组。使用如何根据列拆分 numpy 数组？作为参考，我把这个数组变成了一个数组列表...

df = pd.DataFrame(array)
df['4'] = df[0].astype(str) + df[1].astype(str)
df['4'] = df['4'].astype(int)
arr = df.to_numpy()
y = [arr[arr[:,4]==k] for k in np.unique(arr[:,4])]

y 在哪里...

[array([[0, 0, 1, 2, 0]]),
 array([[0, 1, 0, 0, 1],
        [0, 1, 3, 2, 1],
        [0, 1, 0, 0, 1]]),
 array([[ 1,  1,  0,  1, 11],
        [ 1,  1,  1,  0, 11],
        [ 1,  1,  0,  1, 11]])]

这工作正常，但运行时间太长y。每行花费的时间呈指数增长。我正在玩数亿行，y = [arr[arr[:,4]==k] for k in np.unique(arr[:,4])]从时间的角度来看是不切实际的。

关于如何加快速度的任何想法？

标签： pythonarraysnumpy

解决方案

numpy_indexed使用图书馆怎么样：

import numpy as np
import numpy_indexed as npi

a = np.array([[1, 1, 0, 1],
       [0, 1, 0, 0],
       [1, 1, 1, 0],
       [0, 0, 1, 2],
       [0, 1, 3, 2],
       [1, 1, 0, 1],
       [0, 1, 0, 0]])

key = np.dot(a[:,:2], [1, 10])
y = npi.group_by(key).split_array_as_list(arr)

输出

y
[array([[0, 0, 1, 2]]), 
 array([[0, 1, 0, 0],
        [0, 1, 3, 2],
        [0, 1, 0, 0]]),
 array([[ 1,  1,  0,  1],
        [ 1,  1,  1,  0],
        [ 1,  1,  0,  1]])]

您可以通过以下方式轻松安装库：

> pip install numpy-indexed

推荐阅读