Home > Software design >  matrix with absolute difference between row and column index
matrix with absolute difference between row and column index

Time:01-20

I want to create efficiently matrices of the following shape, i.e. absolute difference between the row and the column number. E.g. for 5 I should get:

0 1 2 3 4
1 0 1 2 3
2 1 0 1 2
3 2 1 0 1
4 3 2 1 0

Is there a fast / elegant way to do this in Python?

CodePudding user response:

You can use a simple nested list comprehension:

n = 5
[[abs(i-j) for j in range(n)]
 for i in range(n)]

Output:

[[0, 1, 2, 3, 4],
 [1, 0, 1, 2, 3],
 [2, 1, 0, 1, 2],
 [3, 2, 1, 0, 1],
 [4, 3, 2, 1, 0]]

Or using numpy:

import numpy as np
n = 5
a = np.arange(n)
abs(a-a[:,None])

CodePudding user response:

Three brilliants more:

np.abs(np.diff(np.indices((n, n)), axis=0))
np.abs(n - 1 - np.sum(np.indices((n, n)), axis=0)[::-1])
np.abs(np.subtract.outer(np.arange(n), np.arange(n)))

I think it's getting a little bit tedious with usage of np.abs and np.indices. Trying to do better...

np.lib.stride_tricks.sliding_window_view(np.cumsum(np.repeat([n-1, -1, 1], [1, n-1, n-1])), n)[::-1].copy()

Done! 4x faster vs 1st example and 2.3x vs 3rd one.

  •  Tags:  
  • Related