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?
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])
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.