Numpy：获取一维数组元素的索引作为二维数组

Question

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

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

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

Answer 1

这是使用 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-only的解决方案：

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

Answer 2

根据您的数据大小，一种可能的选择是退出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.7 秒，这似乎足够合理。

Answer 3

虽然请求的是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 毫秒内处理了 1000 万个项目列表，与用纯 Python 编写的基于列表的版本相比，速度提高了近 50 倍。

对于速度较慢但可读性更强的版本，以下是我之前的版本，基于最近添加的对动态大小的“类型列表”的实验支持，这使我们能够更快地以无序方式填充每个 bin。

这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.

Answer 4

这是一种非常奇怪的方法来做到这一点很糟糕，但我发现不分享太有趣了 - 以及所有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)