CODEWITHSUNDEEP
pythonnumpyfor-looplist-comprehensionnumpy-ndarray
Schneekind
Schneekind

Reputation: 29

What's a pythonic way to generate an array with a diagonal edge?

I created an array with a diagonal edge by using two for loops, but I wonder if there is a simpler way to do so, for example with a list comprehension:

im_diag = np.zeros((im_size, im_size), dtype=np.int8)

for x in range(im_size):
    for y in range(im_size):
        if x+y >= im_size:
            im_diag[x,y] = 1

Output (im_size = 5):

>>> im_size
array([[0, 0, 0, 0, 0],
       [0, 0, 0, 0, 1],
       [0, 0, 0, 1, 1],
       [0, 0, 1, 1, 1],
       [0, 1, 1, 1, 1]], dtype=int8)

Upvotes: 2

Views: 166

Answers (3)

CDJB
CDJB

Reputation: 14506

In general when using numpy, it's far better to use a vectorised method, which will become a lot faster than methods using loops and list-comprehensions as the size of the desired array increases.

>>> np.flip(np.tril(np.ones((5,5)), k=-1), 1)
array([[0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 1.],
       [0., 0., 0., 1., 1.],
       [0., 0., 1., 1., 1.],
       [0., 1., 1., 1., 1.]])

np.ones to create an array of 1s,

np.tril to create a lower triangular array

np.flip to horizontally flip the array

Time Comparison with other answers:

enter image description here

As shown, with n>=10, the numpy method is faster, and as n gets larger, this method will tend to be about 10 times faster than list comprehension solutions.

Code to reproduce:

import perfplot
import numpy as np

def cdjb(n):
    return np.flip(np.tril(np.ones((n,n)), k=-1), 1)

def displayname(n):
    return [[int(j > i) for j in range(n)] for i in range(n)][::-1]

def MikeMajara(n):
    return [[1 if i+j >= n else 0 for i in range(n)] for j in range(n)]

perfplot.show(
    setup=lambda n: n,
    n_range=[2**k for k in range(14)],
    kernels=[
        cdjb,displayname,MikeMajara
        ],
    xlabel='Size of Array',
    )

Upvotes: 5

MikeMajara
MikeMajara

Reputation: 962

If you were looking for a comprehension list this could be your solution:

[[1 if i+j >= im_size else 0 for i in range(im_size)] for j in range(im_size)]

But complex comprehension lists are harder to read/understand. And unreadable code is not pythonic. So give it a thought.

Upvotes: 1

Stefan Falk
Stefan Falk

Reputation: 25387

With list comprehension (no numpy required):

n = 5
[[int(j > i) for j in range(n)] for i in range(n)][::-1]

Gives

[[0, 0, 0, 0, 0],
 [0, 0, 0, 0, 1],
 [0, 0, 0, 1, 1],
 [0, 0, 1, 1, 1],
 [0, 1, 1, 1, 1]]

where [::-1] just reverses the outer list.

Upvotes: 2

Related Questions