CODEWITHSUNDEEP
pythonnumpy
V.Hunon
V.Hunon

Reputation: 320

How to fill a matrix diagonal wise?

I am having trouble of filling a square matrix of n x n with a list of the size m.

What I want to achieve is something like this:

arr = np.array([
    [1, 1, 1, 1],
    [1, 1, 1, 1],
    [1, 1, 1, 1],
    [1, 1, 1, 1]], dtype=float)

And a list like this:

myList= [1, 2, 3, 4, 5, 6, 7]

The output should be:

arr = np.array([
    [4, 5, 6, 7],
    [3, 4, 5, 6],
    [2, 3, 4, 5],
    [1, 2, 3, 4]], dtype=float)

My function works, however I think this is not the most elegant way.

def fill_array(img, myList):

    arr = np.ones(np.shape(img))
    diag_size = (np.shape(img)[0] * 2) - 1
    diag_idx = np.median(np.linspace(1, diag_size, diag_size))

    # First iterate through the lower half, main diagonal and then upper half
    i = 1
    while i <= diag_size:
        factor = myList[i - 1]

        if i < diag_idx:
            position = int(diag_idx - i)
            np.fill_diagonal(arr[position:, :], factor)

        elif i == diag_idx:
            np.fill_diagonal(arr[0:, :], factor)

        elif i > diag_idx:
            position = int(i - diag_idx)
            np.fill_diagonal(arr[:, position:], factor)

        i += 1

    return arr

Is there a better solution for that? Thanks in advance!

Upvotes: 2

Views: 628

Answers (2)

Valdi_Bo
Valdi_Bo

Reputation: 30991

First create a function returning indices of the given diagonal:

def kthDiagIndices(a, offs):
    rows, cols = np.diag_indices_from(a)
    if offs < 0:
        return rows[-offs:], cols[:offs]
    elif offs > 0:
        return rows[:-offs], cols[offs:]
    else:
        return rows, cols

Then define the function to generate the array as:

def fill_array(img, myList):
    arr = np.ones_like(img, dtype=int)
    # How many diagonals can be filled
    diagNo = min(len(myList), img.shape[0] * 2 - 1)
    # Shift from the index in myList to the diagonal offset
    diagShift = diagNo // 2
    # Fill each diagonal in arr
    for i, v in enumerate(myList):
        arr[kthDiagIndices(arr, i - diagShift)] = v
    return arr

Note that this function is "resistant" to passing too long myList (longer than the number of available diagonals).

Another detail to change is to remove dtype=int, but as I saw, you tested your function on integer values.

To test this function, I created:

img = np.ones((4, 4), dtype=int)
myList = [11, 12, 13, 14, 15, 16, 17, 18]

Then I ran:

fill_array(img, myList)

getting:

array([[14, 15, 16, 17],
       [13, 14, 15, 16],
       [12, 13, 14, 15],
       [11, 12, 13, 14]])

Experiment with shorter myLists.

Upvotes: 0

Quang Hoang
Quang Hoang

Reputation: 150765

Let's try as_strided:

from numpy.lib.stride_tricks import as_strided

out = as_strided(np.array(myList, dtype=arr.dtype), 
                 shape=arr.shape, 
                 strides=arr.strides[1:]*2)[::-1]

Output:

array([[4., 5., 6., 7.],
       [3., 4., 5., 6.],
       [2., 3., 4., 5.],
       [1., 2., 3., 4.]])

Upvotes: 3

Related Questions