如何使用numpy有效地求和矩阵的左上角?

时间:2014-02-12 14:40:13

标签: python numpy

天真的蟒蛇方式是:

 sum(grid[x,y] for x in xrange(a) for y in xrange(a-x))

我正在寻找加速,因为grid是numpy中的2d方阵。 tril不起作用,因为它适用于右上角。

2 个答案:

答案 0 :(得分:2)

您可以先旋转矩阵:

np.sum(np.tril(np.rot90(grid)))

请注意,rot90会生成一个视图,而不是数据的副本。下面是几个时序测试,显示两个版本的两个数组大小的相对速度:

In [1]: grid = np.random.rand(10000,10000)

In [2]: a = grid.shape[0]

In [3]: %time sum1 = np.sum(grid[x,y] for x in xrange(a) for y in xrange(a-x))
CPU times: user 18.68 s, sys: 0.04 s, total: 18.72 s
Wall time: 18.60 s

In [4]: %time sum2 = np.sum(np.tril(np.rot90(grid)))
CPU times: user 1.73 s, sys: 0.55 s, total: 2.27 s
Wall time: 2.28 s

第二个版本的速度提高了8倍多。只是为了确认版本是等价的:

In [5]: np.allclose(sum1, sum2)
Out[5]: True

对于较小的阵列:

In [6]: grid = np.random.rand(100,100)

In [7]: a = grid.shape[0]

In [8]: %timeit sum1 = np.sum(grid[x,y] for x in xrange(a) for y in xrange(a-x))
1000 loops, best of 3: 1.9 ms per loop

In [9]: %timeit sum2 = np.sum(np.tril(np.rot90(grid)))
10000 loops, best of 3: 90.4 us per loop

对于这个较小的阵列,第二个版本的速度提高了20多倍。

答案 1 :(得分:0)

@ bogatron的答案稍微冗长一点,可以解释下面发生的事情:

import numpy as np
A = np.arange(25).reshape(5,5)

mask = (np.tri(*A.shape,dtype=bool).T)[:,::-1]

print A[mask].sum()

在这种情况下,面具看起来像:

[[ True  True  True  True  True]
 [ True  True  True  True False]
 [ True  True  True False False]
 [ True  True False False False]
 [ True False False False False]]

这是矩阵的上角。