numpy:将1d数组的元素的索引作为2d数组获取


10

我有一个像这样的numpy数组: [1 2 2 0 0 1 3 5]

是否可以将元素的索引获取为二维数组?例如,上述输入的答案将是[[3 4], [0 5], [1 2], [6], [], [7]]

当前,我必须循环不同的值并numpy.where(input == i)为每个值调用,这在输入足够大的情况下性能很差。


np.argsort([1, 2, 2, 0, 0, 1, 3, 5])array([3, 4, 0, 5, 1, 2, 6, 7], dtype=int64)。那么您可以比较下一个元素。
vb_rises

Answers:


11

这是使用以下方法的O(max(x)+ len(x))方法scipy.sparse

import numpy as np
from scipy import sparse

x = np.array("1 2 2 0 0 1 3 5".split(),int)
x
# array([1, 2, 2, 0, 0, 1, 3, 5])


M,N = x.max()+1,x.size
sparse.csc_matrix((x,x,np.arange(N+1)),(M,N)).tolil().rows.tolist()
# [[3, 4], [0, 5], [1, 2], [6], [], [7]]

通过创建一个稀疏矩阵,其位置为(x [0],0),(x [1],1),...的工作原理CSC很简单。使用(压缩稀疏列)格式。然后将矩阵转换为LIL(链接列表)格式。这种格式将每一行的列索引作为列表存储在其rows属性中,因此我们所需要做的就是将其转换为列表。

请注意,对于小型阵列,argsort基于解决方案的解决方案可能更快,但在某些情况下,这种解决方案可能并非如此。

编辑:

argsort基于numpy解决方案的解决方案:

np.split(x.argsort(kind="stable"),np.bincount(x)[:-1].cumsum())
# [array([3, 4]), array([0, 5]), array([1, 2]), array([6]), array([], dtype=int64), array([7])]

如果组中索引的顺序无关紧要,您也可以尝试argpartition(在这个小例子中碰巧没有什么区别,但这通常不能保证):

bb = np.bincount(x)[:-1].cumsum()
np.split(x.argpartition(bb),bb)
# [array([3, 4]), array([0, 5]), array([1, 2]), array([6]), array([], dtype=int64), array([7])]

编辑:

@Divakar建议不要使用np.split。相反,循环可能更快:

A = x.argsort(kind="stable")
B = np.bincount(x+1).cumsum()
[A[B[i-1]:B[i]] for i in range(1,len(B))]

或者,您可以使用全新的(Python3.8 +)海象运算符:

A = x.argsort(kind="stable")
B = np.bincount(x)
L = 0
[A[L:(L:=L+b)] for b in B.tolist()]

编辑(已编辑):

(不是纯粹的numpy):作为numba的替代方法(请参阅@senderle的文章),我们也可以使用pythran。

编译 pythran -O3 <filename.py>

import numpy as np

#pythran export sort_to_bins(int[:],int)

def sort_to_bins(idx, mx):
    if mx==-1: 
        mx = idx.max() + 1
    cnts = np.zeros(mx + 2, int)
    for i in range(idx.size):
        cnts[idx[i] + 2] += 1
    for i in range(3, cnts.size):
        cnts[i] += cnts[i-1]
    res = np.empty_like(idx)
    for i in range(idx.size):
        res[cnts[idx[i]+1]] = i
        cnts[idx[i]+1] += 1
    return [res[cnts[i]:cnts[i+1]] for i in range(mx)]

在这里numba,以明智的表现取胜:

repeat(lambda:enum_bins_numba_buffer(x),number=10)
# [0.6235917090671137, 0.6071486569708213, 0.6096088469494134]
repeat(lambda:sort_to_bins(x,-1),number=10)
# [0.6235359431011602, 0.6264424560358748, 0.6217901279451326]

较旧的内容:

import numpy as np

#pythran export bincollect(int[:])

def bincollect(a):
    o = [[] for _ in range(a.max()+1)]
    for i,j in enumerate(a):
        o[j].append(i)
    return o

时间与numba(旧)

timeit(lambda:bincollect(x),number=10)
# 3.5732191529823467
timeit(lambda:enumerate_bins(x),number=10)
# 6.7462647299980745

最终比
@Randy

基于循环的方法应该比更好np.split
Divakar

@Divakar好点,谢谢!
Paul Panzer

8

根据数据大小,一种可能的选择是退出numpy并使用collections.defaultdict

In [248]: from collections import defaultdict

In [249]: d = defaultdict(list)

In [250]: l = np.random.randint(0, 100, 100000)

In [251]: %%timeit
     ...: for k, v in enumerate(l):
     ...:     d[v].append(k)
     ...:
10 loops, best of 3: 22.8 ms per loop

然后,您得到的字典{value1: [index1, index2, ...], value2: [index3, index4, ...]}。时间缩放与数组的大小非常接近线性关系,因此10,000,000在我的机器上花费了约2.7s,这似乎足够合理。


7

尽管要求提供numpy解决方案,但我还是决定看看是否有一个有趣的numba基于解决方案的解决方案。确实有!这是一种将分区列表表示为存储在单个预分配缓冲区中的参差不齐的数组的方法。这从Paul Panzerargsort提出的方法中获得了一些启发。(对于性能较差但较简单的旧版本,请参见下文。)

@numba.jit(numba.void(numba.int64[:], 
                      numba.int64[:], 
                      numba.int64[:]), 
           nopython=True)
def enum_bins_numba_buffer_inner(ints, bins, starts):
    for x in range(len(ints)):
        i = ints[x]
        bins[starts[i]] = x
        starts[i] += 1

@numba.jit(nopython=False)  # Not 100% sure this does anything...
def enum_bins_numba_buffer(ints):
    ends = np.bincount(ints).cumsum()
    starts = np.empty(ends.shape, dtype=np.int64)
    starts[1:] = ends[:-1]
    starts[0] = 0

    bins = np.empty(ints.shape, dtype=np.int64)
    enum_bins_numba_buffer_inner(ints, bins, starts)

    starts[1:] = ends[:-1]
    starts[0] = 0
    return [bins[s:e] for s, e in zip(starts, ends)]

这样可以在75毫秒内处理一千万个项目列表,这比使用纯Python编写的基于列表的版本的速度提高了近50倍。

对于较慢但可读性更高的版本,这是我以前的工作,它是基于最近添加的对动态大小的“类型列表”的实验性支持的,这使我们能够以更快的速度无序填充每个垃圾箱。

这与numba的类型推断引擎有些冲突,我敢肯定有更好的方法来处理该部分。事实证明,这比上述方法要慢近10倍。

@numba.jit(nopython=True)
def enum_bins_numba(ints):
    bins = numba.typed.List()
    for i in range(ints.max() + 1):
        inner = numba.typed.List()
        inner.append(0)  # An awkward way of forcing type inference.
        inner.pop()
        bins.append(inner)

    for x, i in enumerate(ints):
        bins[i].append(x)

    return bins

我针对以下各项进行了测试:

def enum_bins_dict(ints):
    enum_bins = defaultdict(list)
    for k, v in enumerate(ints):
        enum_bins[v].append(k)
    return enum_bins

def enum_bins_list(ints):
    enum_bins = [[] for i in range(ints.max() + 1)]
    for x, i in enumerate(ints):
        enum_bins[i].append(x)
    return enum_bins

def enum_bins_sparse(ints):
    M, N = ints.max() + 1, ints.size
    return sparse.csc_matrix((ints, ints, np.arange(N + 1)),
                             (M, N)).tolil().rows.tolist()

我还针对类似于enum_bins_numba_buffer(下面详细介绍)的预编译cython版本对它们进行了测试。

在一千万个随机整数(ints = np.random.randint(0, 100, 10000000))的列表中,我得到以下结果:

enum_bins_dict(ints)
3.71 s ± 80.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

enum_bins_list(ints)
3.28 s ± 52.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

enum_bins_sparse(ints)
1.02 s ± 34.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

enum_bins_numba(ints)
693 ms ± 5.81 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

enum_bins_cython(ints)
82.3 ms ± 1.77 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

enum_bins_numba_buffer(ints)
77.4 ms ± 2.06 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

令人印象深刻的是,即使关闭了边界检查,使用这种方法的numba性能也优于cython相同功能的版本。我还不太熟悉使用pythran这种方法来测试这种方法,但是我很想看到一个比较。基于这种提速,pythran使用此方法的版本似乎也可能更快。

这是cython供参考的版本,并带有一些构建说明。一旦你已经cython安装了,你需要一个简单的setup.py文件是这样的:

from distutils.core import setup
from distutils.extension import Extension
from Cython.Build import cythonize
import numpy

ext_modules = [
    Extension(
        'enum_bins_cython',
        ['enum_bins_cython.pyx'],
    )
]

setup(
    ext_modules=cythonize(ext_modules),
    include_dirs=[numpy.get_include()]
)

和cython模块enum_bins_cython.pyx

# cython: language_level=3

import cython
import numpy
cimport numpy

@cython.boundscheck(False)
@cython.cdivision(True)
@cython.wraparound(False)
cdef void enum_bins_inner(long[:] ints, long[:] bins, long[:] starts) nogil:
    cdef long i, x
    for x in range(len(ints)):
        i = ints[x]
        bins[starts[i]] = x
        starts[i] = starts[i] + 1

def enum_bins_cython(ints):
    assert (ints >= 0).all()
    # There might be a way to avoid storing two offset arrays and
    # save memory, but `enum_bins_inner` modifies the input, and
    # having separate lists of starts and ends is convenient for
    # the final partition stage.
    ends = numpy.bincount(ints).cumsum()
    starts = numpy.empty(ends.shape, dtype=numpy.int64)
    starts[1:] = ends[:-1]
    starts[0] = 0

    bins = numpy.empty(ints.shape, dtype=numpy.int64)
    enum_bins_inner(ints, bins, starts)

    starts[1:] = ends[:-1]
    starts[0] = 0
    return [bins[s:e] for s, e in zip(starts, ends)]

将这两个文件放在您的工作目录中,运行以下命令:

python setup.py build_ext --inplace

然后,您可以使用导入函数from enum_bins_cython import enum_bins_cython


我想知道您是否知道pythran,从广义上讲它类似于numba。我在帖子中添加了pythran解决方案。在这种情况下,pythran似乎占据了上风,提供了一个更快,更Python化的解决方案。
Paul Panzer

@PaulPanzer有趣!我没听说过 我认为,一旦List代码稳定后,numba开发人员将添加预期的语法糖。在这里似乎也有便利/速度之间的权衡-与需要单独的预编译模块的方法相比,jit装饰器很容易集成到普通的Python代码库中。但是,采用scipy方法的3倍加速确实令人印象深刻,甚至令人惊讶!
senderle

只是记得我之前已经完成了此操作:stackoverflow.com/q/55226662/7207392。您介意将numba和cython版本添加到该问答中吗?唯一的区别是:我们不对索引0、1、2,...进行分箱,而是对另一个数组进行分箱。而且,我们不必理会实际切碎的结果数组。
Paul Panzer

@PaulPanzer啊,很酷。我会在今天或明天的某个时候尝试添加它。您是建议一个单独的答案还是对答案进行编辑?无论哪种方式都快乐!
senderle

大!我认为单独发布会更好,但是没有强烈的偏爱。
Paul Panzer

6

这是一种非常怪异的方法,这太糟糕了,但是我发现它实在太有趣了,无法共享-所有人numpy

out = np.array([''] * (x.max() + 1), dtype = object)
np.add.at(out, x, ["{} ".format(i) for i in range(x.size)])
[[int(i) for i in o.split()] for o in out]

Out[]:
[[3, 4], [0, 5], [1, 2], [6], [], [7]]

编辑:这是我沿着这条路可以找到的最佳方法。它仍然比@PaulPanzer的argsort解决方案慢10倍:

out = np.empty((x.max() + 1), dtype = object)
out[:] = [[]] * (x.max() + 1)
coords = np.empty(x.size, dtype = object)
coords[:] = [[i] for i in range(x.size)]
np.add.at(out, x, coords)
list(out)

2

您可以通过制作数字字典来做到这一点,键将是数字,值应是该数字看到的索引,这是最快的方法之一,您可以看到以下代码:

>>> import numpy as np
>>> a = np.array([1 ,2 ,2 ,0 ,0 ,1 ,3, 5])
>>> b = {}
# Creating an empty list for the numbers that exist in array a
>>> for i in range(np.min(a),np.max(a)+1):
    b[str(i)] = []

# Adding indices to the corresponding key
>>> for i in range(len(a)):
    b[str(a[i])].append(i)

# Resulting Dictionary
>>> b
{'0': [3, 4], '1': [0, 5], '2': [1, 2], '3': [6], '4': [], '5': [7]}

# Printing the result in the way you wanted.
>>> for i in sorted (b.keys()) :
     print(b[i], end = " ")

[3, 4] [0, 5] [1, 2] [6] [] [7] 

1

伪代码:

  1. 通过从最大值中减去numpy数组的最小值,然后加1,得到“ 2d数组中1d数组的数量”。在您的情况下,它将是5-0 + 1 = 6

  2. 使用其中的1d数组数量初始化2d数组。在您的情况下,请初始化其中包含6个1d数组的2d数组。每个1d数组对应于numpy数组中的唯一元素,例如,第一个1d数组将对应于“ 0”,第二个1d数组将对应于“ 1”,...

  3. 遍历您的numpy数组,将元素的索引放入正确的对应1d数组中。在您的情况下,numpy数组中第一个元素的索引将放入第二个1d数组,numpy数组中第二个元素的索引将放入第三个1d数组,...。

此伪代码将花费线性时间来运行,因为它取决于numpy数组的长度。


1

这将为您提供所需的确切信息,并且在我的计算机上花费10,000,000大约需要2.5秒:

import numpy as np
import timeit

# x = np.array("1 2 2 0 0 1 3 5".split(),int)
x = np.random.randint(0, 100, 100000)

def create_index_list(x):
    d = {}
    max_value = -1
    for i,v in enumerate(x):
        if v > max_value:
            max_value = v
        try:
            d[v].append(i)
        except:
            d[v] = [i]
    result_list = []
    for i in range(max_value+1):
        if i in d:
            result_list.append(d[i])
        else:
            result_list.append([])
    return result_list

# print(create_index_list(x))
print(timeit.timeit(stmt='create_index_list(x)', number=1, globals=globals()))


0

因此,给定一个元素列表,您希望建立(元素,索引)对。在线性时间内,可以这样执行:

hashtable = dict()
for idx, val in enumerate(mylist):
    if val not in hashtable.keys():
         hashtable[val] = list()
    hashtable[val].append(idx)
newlist = sorted(hashtable.values())

这应该花费O(n)时间。到目前为止,我还没有想到一个更快的解决方案,但是如果可以的话,将在这里进行更新。

By using our site, you acknowledge that you have read and understand our Cookie Policy and Privacy Policy.
Licensed under cc by-sa 3.0 with attribution required.