Vectorized solution for generating NumPy matrix

Question

I am looking for some built-in NumPy module or some vectorized approach to get such n by n matrices for n>1. The only key thing here is that the last element of a given row serves as the first element of the following row.

n = 2
# array([[1, 2], 
#        [2, 3]])

n = 3
# array([[1, 2, 3], 
#        [3, 4, 5], 
#        [5, 6, 7]])

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

My attempt using list comprehensions. The same can be written in an extended for loop syntax as well.

import numpy as np
n = 4
arr = np.array([[(n-1)*j+i for i in range(1, n+1)] for j in range(n)])
# array([[ 1,  2,  3,  4],
#        [ 4,  5,  6,  7],
#        [ 7,  8,  9, 10],
#        [10, 11, 12, 13]])

Answer 1

If you want to keep things simple (and somewhat readable), this should do it:

def ranged_mat(n):
    out = np.arange(1, n ** 2 + 1).reshape(n, n)
    out -= np.arange(n).reshape(n, 1)
    return out

Simply build all numbers from 1 to n², reshape them to the desired block shape, then subtract the row number from each row.

This does the same as Divakar's ranged_mat_v2, but I like being explicit with intermediate array shapes. Not everyone is an expert at NumPy's broadcasting rules.

Answer 2

Approach #1

We can leverage np.lib.stride_tricks.as_strided based scikit-image's view_as_windows to get sliding windows. More info on use of as_strided based view_as_windows.

Additionally, it accepts a step argument and that would fit in perfectly for this problem. Hence, the implementation would be -

from skimage.util.shape import view_as_windows

def ranged_mat(n):
    r  = np.arange(1,n*(n-1)+2)
    return view_as_windows(r,n,step=n-1)

Sample runs -

In [270]: ranged_mat(2)
Out[270]: 
array([[1, 2],
       [2, 3]])

In [271]: ranged_mat(3)
Out[271]: 
array([[1, 2, 3],
       [3, 4, 5],
       [5, 6, 7]])

In [272]: ranged_mat(4)
Out[272]: 
array([[ 1,  2,  3,  4],
       [ 4,  5,  6,  7],
       [ 7,  8,  9, 10],
       [10, 11, 12, 13]])

Approach #2

Another with outer-broadcasted-addition -

def ranged_mat_v2(n):
    r = np.arange(n)
    return (n-1)*r[:,None]+r+1

Approach #3

We can also use numexpr module that supports multi-core processing and hence achieve better efficiency on large n's -

import numexpr as ne

def ranged_mat_v3(n):
    r = np.arange(n)
    r2d = (n-1)*r[:,None]
    return ne.evaluate('r2d+r+1')

Making use of slicing gives us a more memory-efficient one -

def ranged_mat_v4(n):
    r  = np.arange(n+1)
    r0 = r[1:]
    r1 = r[:-1,None]*(n-1)
    return ne.evaluate('r0+r1')

---

Timings -

In [423]: %timeit ranged_mat(10000)
273 ms ± 3.42 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [424]: %timeit ranged_mat_v2(10000)
316 ms ± 2.03 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [425]: %timeit ranged_mat_v3(10000)
176 ms ± 85.9 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [426]: %timeit ranged_mat_v4(10000)
154 ms ± 82.8 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Answer 3

There is also np.fromfunction as below. docs here

def func(n):
    return np.fromfunction(lambda r,c: (n-1)*r+1+c, shape=(n,n))

It takes a function which calculates an array from the index values.

Answer 4

We can use NumPy strides for this:

def as_strides(n):
    m = n**2 - (n-1)
    a = np.arange(1, m+1)
    s = a.strides[0]
    return np.lib.stride_tricks.as_strided(a, shape=(n,n), strides=((n-1)*s,s))

---

as_strides(2)
rray([[1, 2],
       [2, 3]])

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

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