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

Question

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

我尝试使用 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) 我想获得以下输出：

   0  1  2  3  4  5  6  7
0  O  O  O  O  O  X  O  O
1  O  O  O  O  X  O  O  O
2  O  O  O  X  O  O  O  O
3  X  O  X  O  O  O  O  O
4  O  X  O  O  O  O  O  O
5  X  O  X  O  O  O  O  O
6  O  O  O  X  O  O  O  O
7  O  O  O  O  X  O  O  O

Answer 1

方法 #1：Broadcasted-masking

这是 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。这个想法是找出对角线和 [=36=]anti-diagonal 的起始扁平化索引，并在扁平化的 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)

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

Fill diagonal and anti-diagonal offsetted by row, column indices in array - NumPy / Python

python

numpy

numpy-ndarray