CODEWITHSUNDEEP
pythonarraysnumpy
Viktor
Viktor

Reputation: 601

Create 2D Matrix of ascending integers in diagonal/triangle-like order with Numpy

How do I create a matrix of ascending integers that are arrayed like this example of N=6?

1  3  6
2  5  0 
4  0  0

Here another example for N=13:

1  3  6  10 0
2  5  9  13 0
4  8  12 0  0
7  11 0  0  0
10 0  0  0  0

Also, the solution should perform well for large N values.

My code

import numpy as np
N = 13
array_dimension = 5
x = 0
y = 1
z = np.zeros((array_dimension,array_dimension))
z[0][0] = 1
for i in range(2, N+1):
    z[y][x] = i
    if y == 0:
        y = (x + 1)
        x = 0
    else:
        x += 1
        y -= 1
print(z)

[[ 1.  3.  6. 10.  0.]
 [ 2.  5.  9.  0.  0.]
 [ 4.  8. 13.  0.  0.]
 [ 7. 12.  0.  0.  0.]
 [11.  0.  0.  0.  0.]]

works, but there must be a more efficient way. Most likely via Numpy, but I cannot find a solution.

Upvotes: 6

Views: 331

Answers (2)

Mechanic Pig
Mechanic Pig

Reputation: 7736

The assignment can be completed in one step by simply transforming the index of the lower triangle:

def fill_diagonal(n):
    assert n > 0
    m = int((2 * n - 1.75) ** 0.5 + 0.5)
    '''n >= ((1 + (m - 1)) * (m - 1)) / 2 + 1
    => 2n - 2 >= m ** 2 - m
    => 2n - 7 / 4 >= (m - 1 / 2) ** 2
    => (2n - 7 / 4) ** (1 / 2) + 1 / 2 >= m for n > 0
    => m = floor((2n - 7 / 4) ** (1 / 2) + 1 / 2)
    or n <= ((1 + m) * m) / 2
    => (2n + 1 / 4) ** (1 / 2) - 1 / 2 <= m for n > 0
    => m = ceil((2n + 1 / 4) ** (1 / 2) - 1 / 2)
    '''
    i, j = np.tril_indices(m)
    i -= j
    ret = np.zeros((m, m), int)
    ret[i[:n], j[:n]] = np.arange(1, n + 1)
    return ret

Test:

>>> for i in range(1, 16):
...     print(repr(fill_diagonal(i)), end='\n\n')
...
array([[1]])

array([[1, 0],
       [2, 0]])

array([[1, 3],
       [2, 0]])

array([[1, 3, 0],
       [2, 0, 0],
       [4, 0, 0]])

array([[1, 3, 0],
       [2, 5, 0],
       [4, 0, 0]])

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

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

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

array([[1, 3, 6, 0],
       [2, 5, 9, 0],
       [4, 8, 0, 0],
       [7, 0, 0, 0]])

array([[ 1,  3,  6, 10],
       [ 2,  5,  9,  0],
       [ 4,  8,  0,  0],
       [ 7,  0,  0,  0]])

array([[ 1,  3,  6, 10,  0],
       [ 2,  5,  9,  0,  0],
       [ 4,  8,  0,  0,  0],
       [ 7,  0,  0,  0,  0],
       [11,  0,  0,  0,  0]])

array([[ 1,  3,  6, 10,  0],
       [ 2,  5,  9,  0,  0],
       [ 4,  8,  0,  0,  0],
       [ 7, 12,  0,  0,  0],
       [11,  0,  0,  0,  0]])

array([[ 1,  3,  6, 10,  0],
       [ 2,  5,  9,  0,  0],
       [ 4,  8, 13,  0,  0],
       [ 7, 12,  0,  0,  0],
       [11,  0,  0,  0,  0]])

array([[ 1,  3,  6, 10,  0],
       [ 2,  5,  9, 14,  0],
       [ 4,  8, 13,  0,  0],
       [ 7, 12,  0,  0,  0],
       [11,  0,  0,  0,  0]])

array([[ 1,  3,  6, 10, 15],
       [ 2,  5,  9, 14,  0],
       [ 4,  8, 13,  0,  0],
       [ 7, 12,  0,  0,  0],
       [11,  0,  0,  0,  0]])

For the case of large n, the performance is about 10 to 20 times that of the loop solution:

%timeit fill_diagonal(10 ** 5)
1.63 ms ± 94.5 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
%timeit fill_diagonal_loop(10 ** 5)    # OP's solution
25.1 ms ± 218 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Upvotes: 4

Ali_Sh
Ali_Sh

Reputation: 2816

I didn't go into the depth of your code and try to write a better code, but just accelerating your code with numba library decorator can improve its efficiency much than twenty times (even hundreds):

import numpy as np
import numba as nb

@nb.njit
def new_(N, array_dimension):
    x = 0
    y = 1
    z = np.zeros((array_dimension,array_dimension))
    z[0, 0] = 1
    for i in range(2, N+1):
        z[y, x] = i
        if y == 0:
            y = (x + 1)
            x = 0
        else:
            x += 1
            y -= 1
    return z

Benchmark: (Colab temp link)

enter image description here

It seems numba is the better choice rather than numpy in this regard.

Upvotes: 3

Related Questions