CODEWITHSUNDEEP
pythonnumpynumpy-ndarray
the_dummy
the_dummy

Reputation: 347

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

I have an n by n matrix. Given a position, I would like to replace the diagonals of that position with a specific value.

I tried using np.fill_diagonal but I don't get the correct result.

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")

For update_available(4,1) I want to get the following output: 

   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

Upvotes: 2

Views: 1219

Answers (1)

Divakar
Divakar

Reputation: 221584

Approach #1 : Broadcasted-masking

Here's one with 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

Approach #2 : NumPy-Eye-masking

Putting in those same offsets into np.eye, we could make it a bit more compact, like so -

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

Approach #3: Flattened assignment

We can also make use of slicing. The idea is figure out the starting flattened indices for the diagonal and anti-diagonal ones and slice the array at flattened n+1 (n = length of array) steps and assign the new values. This should be much more efficient, especially on large arrays. The implementation would look something like this -

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

Sample run -

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]])

Timings on large arrays -

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)

Upvotes: 2

Related Questions