填充数组中按行,列索引偏移的对角线和反对角线-NumPy / Python

时间:2019-06-07 04:13:38

标签: python numpy numpy-ndarray

我有一个n x n的矩阵。给定位置,我想用特定值替换该位置的对角线。

我尝试使用ConnectionPool,但没有得到正确的结果。

np.fill_diagonal

对于import numpy as np n = 8 available = np.array([["O" for _ in range(n)] for _ in range(n)]) def update_available(column, row): available[column][row] = "X" np.fill_diagonal(available[column-row:],"X") np.fill_diagonal(available[::-1,:][column-row-1:],"X") ,我想获得以下输出:

update_available(4,1)

1 个答案:

答案 0 :(得分:1)

方法1:广播掩蔽

这里是masking-

def fill_cross(a, row, col, newval):
    n = len(a)
    r = np.arange(n)
    m1 = r[::-1,None] == r + n-row-col-1
    m2 = r[:,None] == r+row-col
    a[m1|m2] = newval
    return a

方法2:用NumPy-Eye-masking

将这些相同的偏移量放入np.eye中,我们可以使它更紧凑 ,就像这样-

def fill_cross_v2(a, row, col, newval):
    n = len(a)
    m1 = np.eye(n,k=-row+col,dtype=bool)
    m2 = np.eye(n,k=n-col-row-1,dtype=bool)[:,::-1]
    a[m1|m2] = newval
    return a

方法3:扁平化分配

我们还可以使用slicing。这个想法是找出对角线和反对角线的起始变平索引,并以变平的n+1(n =数组长度)步长对数组进行切片并分配新值。这应该效率更高,尤其是在大型阵列上。实现看起来像这样-

def fill_cross_v3(a, row, col, newval):
    if row+col>=n:
        anti_diag_start = (row+col-n+1,n-1)
    else:
        anti_diag_start = (0,row+col)

    if row>col:
        diag_start = (row-col,0)
    else:
        diag_start = (0,col-row)

    r,c = [np.ravel_multi_index(i,a.shape) for i in [diag_start,anti_diag_start]]
    a.ravel()[r:r+(n-diag_start[0]-diag_start[1])*(n+1):n+1] = newval
    a.ravel()[c:c*(n+1):n-1] = newval
    return a

样品运行-

In [71]: a
Out[71]: 
array([[0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0]])

In [72]: fill_cross(a, row=4, col=1, newval=-1)
Out[72]: 
array([[ 0,  0,  0,  0,  0, -1,  0,  0],
       [ 0,  0,  0,  0, -1,  0,  0,  0],
       [ 0,  0,  0, -1,  0,  0,  0,  0],
       [-1,  0, -1,  0,  0,  0,  0,  0],
       [ 0, -1,  0,  0,  0,  0,  0,  0],
       [-1,  0, -1,  0,  0,  0,  0,  0],
       [ 0,  0,  0, -1,  0,  0,  0,  0],
       [ 0,  0,  0,  0, -1,  0,  0,  0]])

大型数组上的计时-

In [509]: a = np.zeros((1000,1000))

In [510]: %timeit fill_cross(a, row=200, col=700, newval=-1)
     ...: %timeit fill_cross_v2(a, row=200, col=700, newval=-1)
     ...: %timeit fill_cross_v3(a, row=200, col=700, newval=-1)
1.64 ms ± 15.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
632 µs ± 787 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)
438 µs ± 12.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [511]: a = np.zeros((5000,5000))

In [512]: %timeit fill_cross(a, row=2200, col=2700, newval=-1)
     ...: %timeit fill_cross_v2(a, row=2200, col=2700, newval=-1)
     ...: %timeit fill_cross_v3(a, row=2200, col=2700, newval=-1)
66.6 ms ± 680 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
39.1 ms ± 245 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
17.3 ms ± 50.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)