慕的地8271018
您可以使用以下想法。如果major, minor = A.strides然后将步幅设置A为major + minor, minor（应小心操作以避免超出数组边界），您将获得每列为对角线的数组。通过这种方式，A.sum(axis=0)您可以计算对角线之和。对于意味着您可以使用相同但乘以某些值，例如A.shape[0] / [1, 2, ... A.shape[0], ... 2, 1]修复对角线长度的变化。对于方差，您可以使用它variance = <(x - <x>)**2> = <x**2> - <x>**2。import numpy as npdef rot45(A): """ >>> A = np.triu(np.arange(25).reshape(5, 5), 1) >>> print(A) [[ 0 1 2 3 4] [ 0 0 7 8 9] [ 0 0 0 13 14] [ 0 0 0 0 19] [ 0 0 0 0 0]] >>> print(rot45(A)) [[ 1 2 3 4] [ 7 8 9 0] [13 14 0 0] [19 0 0 0]] """ major, minor = A.strides strides = major + minor, minor shape = A.shape[0] - 1, A.shape[1] return np.lib.stride_tricks.as_strided(A, shape, strides)[:, 1:]def apply_diag(A, func): """ >>> A = np.arange(25).reshape(5, 5) >>> print(A) [[ 0 1 2 3 4] [ 5 6 7 8 9] [10 11 12 13 14] [15 16 17 18 19] [20 21 22 23 24]] >>> offset_list = np.arange(-1 * len(A) + 1, len(A)) >>> diag_var_list = [np.sum(np.diagonal(A, k)) for k in offset_list] >>> diag_var_list [20, 36, 48, 56, 60, 40, 24, 12, 4] >>> print(apply_diag(A, np.sum)) [20, 36, 48, 56, 60, 40, 24, 12, 4] """ U = np.triu(A, 1) U = rot45(U) D = np.tril(A, -1).T.copy() D = rot45(D) return func(D, axis=0)[::-1].tolist() + [func(np.diag(A))] + func(U, axis=0)[::-1].tolist()[::-1]def using_numpy(A): """ >>> A = np.arange(25).reshape(5, 5) >>> print(A) [[ 0 1 2 3 4] [ 5 6 7 8 9] [10 11 12 13 14] [15 16 17 18 19] [20 21 22 23 24]] >>> offset_list = np.arange(-1 * len(A) + 1, len(A)) >>> diag_var_list = [np.var(np.diagonal(A, k)) for k in offset_list] >>> diag_var_list [0.0, 9.0, 24.0, 45.0, 72.0, 45.0, 24.0, 9.0, 0.0] >>> print(using_numpy(A)) [ 0. 9. 24. 45. 72. 45. 24. 9. 0.] """ multiply = (A.shape[0] - 1) / np.r_[1:A.shape[0], A.shape[0] - 1, A.shape[0] - 1:0:-1] return multiply * apply_diag(A ** 2, np.mean) - (multiply * apply_diag(A, np.mean))**2def list_comp(A, func): """ >>> A = np.arange(25).reshape(5, 5) >>> list_comp(A, np.sum) [20, 36, 48, 56, 60, 40, 24, 12, 4] """ offset_list = np.arange(-1 * len(A) + 1, len(A)) return [func(np.diagonal(A, k)) for k in offset_list]对于 (100, 100) 大小的矩阵，速度似乎提高了 10 倍，但对于更大的矩阵，速度差异下降得更低。In [9]: A = np.random.randn(100, 100)In [10]: %timeit using_numpy(A)761 µs ± 3.35 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)In [11]: %timeit list_comp(A, np.var)9.57 ms ± 19.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)In [12]: A = np.random.randn(1000, 1000)In [13]: %timeit using_numpy(A)37.4 ms ± 125 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)In [14]: %timeit list_comp(A, np.var)112 ms ± 927 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)