С помощью numpy заполнить соседние с главной диагонали

Рейтинг: 1Ответов: 3Опубликовано: 27.04.2015

Как с помощью numpy создать массив:

[1 -2   1              ]
|   1  -2  1           |
|          .           |
|             .        |
|                .     |
[              1  -2  1]

ОБНОВЛЕНИЕ

Могу, например, заполнить главную диагональ:

>>> a = np.zeros((3, 3), int)
>>> np.fill_diagonal(a, -2)
>>> a
array([[-2, 0, 0],
       [0, -2, 0],
       [0, 0, -2]])

Затрудняюсь заполнить нужным числом соседние с главной диагонали. Могу решить задачу так:

def ff (i , j):
if i==j:
    return -2
elif i-j == 1 or i-j == -1:
    return 1
else:
    return 0


D = np.array([ff(i, j) for i, j in itertools.product(range(N),
range(N))]).reshape([N, N])

Может кто-нибудь предложить более короткое решение?

python
Источник: Stack Overflow на русском

Ответы

▲ 2Принят

Где k - номер диагонали, начиная от центральной

import numpy as np

array = np.zeros((5, 5))
k = 1

rows, cols = np.indices(array.shape)
row_values = np.diag(rows, k=k)
col_values = np.diag(cols, k=k)
array[row_values, col_values] = 1

print(array)

Вывод программы:

[[ 0.  1.  0.  0.  0.]
 [ 0.  0.  1.  0.  0.]
 [ 0.  0.  0.  1.  0.]
 [ 0.  0.  0.  0.  1.]
 [ 0.  0.  0.  0.  0.]]
▲ 3

Можно сделать короче:

array = np.zeros((5, 5))
k = 1

rows, cols = np.indices(array.shape)
array[rows == cols-k] = 1
print(array)

Вывод программы:

[[ 0.  1.  0.  0.  0.]
 [ 0.  0.  1.  0.  0.]
 [ 0.  0.  0.  1.  0.]
 [ 0.  0.  0.  0.  1.]
 [ 0.  0.  0.  0.  0.]]

Но по результатам тестирования, это решение на больших матрицах будет работать медленнее почти в два раза, по сравнению с решением @Avernial.

Для тестирования скорости были созданы две функции:

def change_diag_1(size, k=1):
    a = np.zeros((size, size))
    rows, cols = np.indices(a.shape)
    r_v = np.diag(rows, k=k)
    c_v = np.diag(cols, k=k)
    a[r_v, c_v] = 10

def change_diag_2(size, k=1):
    a = np.zeros((size, size))
    i, j = np.indices(a.shape)
    a[i == j-k] = 10

Использовалась команда %timeit в ipython.

python версии 3.4.3, numpy версии 1.9.2.

Полученные результаты:

+------+---------------+---------------+
| size | change_diag_1 | change_diag_2 |
+------+---------------+---------------+
|   10 | 40.3 µs*      | 35.8 µs*      |
|  100 | 104 µs*       | 122 µs*       |
| 1000 | 9.85 ms       | 14.4 ms       |
| 5000 | 174 ms        | 302 ms        |
+------+---------------+---------------+

* timeit: разница по времени работы между самым быстрым и самым 
          медленным запуском функции в несколько раз (конкретно 
          на этих запусках от 7 до 11). Это потенциально вызвано
          кэшированием промежуточных результатов.

Полный отчёт о запусках:

> %timeit change_diag_1(10)
The slowest run took 6.94 times longer than the fastest. This could mean that an intermediate result is being cached 
10000 loops, best of 3: 40.3 µs per loop

> %timeit change_diag_2(10)
The slowest run took 8.65 times longer than the fastest. This could mean that an intermediate result is being cached 
10000 loops, best of 3: 35.8 µs per loop


> %timeit change_diag_1(100)
The slowest run took 11.24 times longer than the fastest. This could mean that an intermediate result is being cached 
10000 loops, best of 3: 104 µs per loop

> %timeit change_diag_2(100) 
The slowest run took 8.62 times longer than the fastest. This could mean that an intermediate result is being cached 
10000 loops, best of 3: 122 µs per loop

> %timeit change_diag_1(1000)  
100 loops, best of 3: 9.85 ms per loop

> %timeit change_diag_2(1000) 
100 loops, best of 3: 14.4 ms per loop

> %timeit change_diag_1(5000)
10 loops, best of 3: 174 ms per loop

> %timeit change_diag_2(5000)
1 loops, best of 3: 302 ms per loop
▲ 1

Вот такое решение показало лучший перфоманс по примеру микробенчмарка от Тимофея:

k = 1
size = 5
a = np.zeros((size,size))
x = np.arange(0,size-k)
y = x + k
a[x,y] = 1