从没有中间索引数组的numpy数组中选择的快速方法

Question

给定以下 2 列数组，我想从第二列中选择与第一列中的“边”相对应的项目。这只是一个示例，因为实际上我的a 可能有数百万行。因此，理想情况下，我希望尽可能快地完成此操作，并且不会产生中间结果。

import numpy as np
a = np.array([[1,4],[1,2],[1,3],[2,6],[2,1],[2,8],[2,3],[2,1],
              [3,6],[3,7],[5,4],[5,9],[5,1],[5,3],[5,2],[8,2],
              [8,6],[8,8]])

即我要查找结果，

desired = np.array([4,6,6,4,2])

这是a[:,1] 中的条目，对应于a[:,0] 更改的位置。

一个解决方案是，

b = a[(a[1:,0]-a[:-1,0]).nonzero()[0]+1, 1]

给出np.array([6,6,4,2])，我可以简单地添加第一项，没问题。但是，这会创建第一个项目的索引的中间数组。我可以通过使用列表理解来避免中间：

c = [a[i+1,1] for i,(x,y) in enumerate(zip(a[1:,0],a[:-1,0])) if x!=y]

这也给出了[6,6,4,2]。假设基于生成器的zip（在 Python 3 中为真），这不需要创建中间表示并且应该非常节省内存。但是，内部循环不是 numpy，它需要生成一个列表，然后必须将其转换回 numpy 数组。

你能想出一个内存效率为c但速度效率为b的numpy-only版本吗？理想情况下，只需要通过a 一次。

（请注意，在这里测量速度不会有太大帮助，除非 a 非常大，所以我不会费心对此进行基准测试，我只想要理论上快速且内存高效的东西。例如，你可以假设a 中的行是从文件流式传输的并且访问速度很慢——避免b 解决方案的另一个原因，因为它需要通过a 进行第二次随机访问。）

编辑：一种生成用于测试的大型a 矩阵的方法：

from itertools import repeat
N, M = 100000, 100
a = np.array(zip([x for y in zip(*repeat(np.arange(N),M)) for x in y ], np.random.random(N*M)))

Answer 1

恐怕如果您希望以矢量化方式执行此操作，则无法避免中间数组，因为它没有内置数组。

现在，让我们寻找除nonzero() 之外的向量化方法，它可能会更高效。按照与(a[1:,0]-a[:-1,0]) 的原始代码执行微分相同的想法，我们可以在查找对应于“边缘”或移位的非零微分之后使用布尔索引。

因此，我们将采用像这样的矢量化方法 -

a[np.append(True,np.diff(a[:,0])!=0),1]

运行时测试

原来的解决方案a[(a[1:,0]-a[:-1,0]).nonzero()[0]+1,1] 会跳过第一行。但是，为了计时的目的，我们只是说，这是一个有效的结果。这是针对本文中提出的解决方案的运行时 -

In [118]: from itertools import repeat
     ...: N, M = 100000, 2
     ...: a = np.array(zip([x for y in zip(*repeat(np.arange(N),M))\
                              for x in y ], np.random.random(N*M)))
     ...: 

In [119]: %timeit a[(a[1:,0]-a[:-1,0]).nonzero()[0]+1,1]
100 loops, best of 3: 6.31 ms per loop

In [120]: %timeit a[1:][np.diff(a[:,0])!=0,1]
100 loops, best of 3: 4.51 ms per loop

现在，假设您也想包含第一行。更新的运行时看起来像这样 -

In [123]: from itertools import repeat
     ...: N, M = 100000, 2
     ...: a = np.array(zip([x for y in zip(*repeat(np.arange(N),M))\
                              for x in y ], np.random.random(N*M)))
     ...: 

In [124]: %timeit a[np.append(0,(a[1:,0]-a[:-1,0]).nonzero()[0]+1),1]
100 loops, best of 3: 6.8 ms per loop

In [125]: %timeit a[np.append(True,np.diff(a[:,0])!=0),1]
100 loops, best of 3: 5 ms per loop

Answer 2

好吧其实我找到了解决办法，刚刚学习了np.fromiter，它可以基于一个生成器构建一个numpy数组：

d = np.fromiter((a[i+1,1] for i,(x,y) in enumerate(zip(a[1:,0],a[:-1,0])) if x!=y), int)

我认为这样做会生成一个没有任何中间数组的 numpy 数组。但是，需要注意的是，它似乎并没有那么有效！忘记我在关于测试的问题中所说的：

t = [lambda a: a[(a[1:,0]-a[:-1,0]).nonzero()[0]+1, 1],
     lambda a: np.array([a[i+1,1] for i,(x,y) in enumerate(zip(a[1:,0],a[:-1,0])) if x!=y]),
     lambda a: np.fromiter((a[i+1,1] for i,(x,y) in enumerate(zip(a[1:,0],a[:-1,0])) if x!=y), int)]

from timeit import Timer
[Timer(x(a)).timeit(number=10) for x in t]

[0.16596235800034265, 1.811289312000099, 2.1662971739997374]

似乎第一个解决方案要快得多！我认为这是因为即使它生成中间数据，它也能够在 numpy 中完全执行内部循环，而在另一个中它为数组中的每个项目运行 Python 代码。

就像我说的，这就是为什么我不确定这种基准测试在这里是否有意义——如果对a 的访问速度要慢得多，那么基准测试就不会加载 CPU。想法？

不“接受”这个答案，因为我希望有人能更快地想出一些东西。

Answer 3

如果您关心内存效率，可以这样解决：与输入数据相同大小顺序的唯一中间体可以由 bool 类型（a[1:,0] != a[:- 1, 0]);如果您的输入数据是 int32，则它比 'a' 本身小 8 倍。您也可以计算该二进制数组的非零值以预分配输出数组；尽管如果 != 的输出像您的示例所暗示的那样稀疏，那么这应该不是很重要。