在numpy中查找对角线总和(更快)

本文介绍了在numpy中查找对角线总和(更快)的处理方法，对大家解决问题具有一定的参考价值，需要的朋友们下面随着小编来一起学习吧！ 问题描述

我有一些像这样的board numpy数组:

I have some board numpy arrays like that:

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

我正在使用以下代码在电路板的-7至8(及其镜像版本)的第n个对角线上找到元素的总和.

And I'm using the following code to find the sum of elements on each nth diagonal from -7 to 8 of the board (and the mirrored version of it).

n = 8 rate = [b.diagonal(i).sum() for b in (board, board[::-1]) for i in range(-n+1, n)]

经过分析后，此操作大约花费了总运行时间的2/3，这似乎是由于以下两个因素造成的:

After some profiling, this operation is taking about 2/3 of overall running time and it seems to be because of 2 factors:

• .diagonal方法将构建一个新数组，而不是一个视图(看起来numpy 1.7将具有一个新的.diag方法来解决该问题)
• 迭代是在列表理解内的python中完成的

• The .diagonal method builds a new array instead of a view (looks like numpy 1.7 will have a new .diag method to solve that)
• The iteration is done in python inside the list comprehension

那么，有什么方法可以更快地找到这些总和(可能在numpy的C层中)?

So, there are any methods to find these sums faster (possibly in the C layer of numpy)?

经过更多测试后，通过缓存此操作，我可以减少7.5倍的总时间...也许我在寻找错误的瓶颈?

After some more tests, I could reduce 7.5x the total time by caching this operation... Maybe I was looking for the wrong bottleneck?

还有一件事:

只是找到了.trace方法来代替diagonal(i).sum()东西，并且...性能没有太大提高(大约2％到4％).

Just found the .trace method that replaces the diagonal(i).sum() thing and... There wasn't much improvement in performance (about 2 to 4%).

所以问题应该出在理解上.有什么想法吗?

So the problem should be the comprehension. Any ideas?

推荐答案

有可能使用stride_tricks的解决方案.这部分是基于此问题的答案中可用的大量信息，但是问题恰恰不同，我认为，不应将其视为重复项.这是应用于方阵的基本思想，有关实现更通用解决方案的函数，请参见下文.

There's a possible solution using stride_tricks. This is based in part on the plethora of information available in the answers to this question, but the problem is just different enough, I think, not to count as a duplicate. Here's the basic idea, applied to a square matrix -- see below for a function implementing the more general solution.

>>> cols = 8 >>> a = numpy.arange(cols * cols).reshape((cols, cols)) >>> fill = numpy.zeros((cols - 1) * cols, dtype='i8').reshape((cols - 1, cols)) >>> stacked = numpy.vstack((a, fill, a)) >>> major_stride, minor_stride = stacked.strides >>> strides = major_stride, minor_stride * (cols + 1) >>> shape = (cols * 2 - 1, cols) >>> numpy.lib.stride_tricks.as_strided(stacked, shape, strides) array([[ 0, 9, 18, 27, 36, 45, 54, 63], [ 8, 17, 26, 35, 44, 53, 62, 0], [16, 25, 34, 43, 52, 61, 0, 0], [24, 33, 42, 51, 60, 0, 0, 0], [32, 41, 50, 59, 0, 0, 0, 0], [40, 49, 58, 0, 0, 0, 0, 0], [48, 57, 0, 0, 0, 0, 0, 0], [56, 0, 0, 0, 0, 0, 0, 0], [ 0, 0, 0, 0, 0, 0, 0, 7], [ 0, 0, 0, 0, 0, 0, 6, 15], [ 0, 0, 0, 0, 0, 5, 14, 23], [ 0, 0, 0, 0, 4, 13, 22, 31], [ 0, 0, 0, 3, 12, 21, 30, 39], [ 0, 0, 2, 11, 20, 29, 38, 47], [ 0, 1, 10, 19, 28, 37, 46, 55]]) >>> diags = numpy.lib.stride_tricks.as_strided(stacked, shape, strides) >>> diags.sum(axis=1) array([252, 245, 231, 210, 182, 147, 105, 56, 7, 21, 42, 70, 105, 147, 196])

当然，我不知道这实际上有多快.但是我敢打赌它会比Python列表理解要快.

Of course, I have no idea how fast this will actually be. But I bet it will be faster than a Python list comprehension.

为方便起见，这是一个完全通用的diagonals函数.假设您想沿最长轴移动对角线:

For convenience, here's a fully general diagonals function. It assumes you want to move the diagonal along the longest axis:

def diagonals(a): rows, cols = a.shape if cols > rows: a = a.T rows, cols = a.shape fill = numpy.zeros(((cols - 1), cols), dtype=a.dtype) stacked = numpy.vstack((a, fill, a)) major_stride, minor_stride = stacked.strides strides = major_stride, minor_stride * (cols + 1) shape = (rows + cols - 1, cols) return numpy.lib.stride_tricks.as_strided(stacked, shape, strides)