3

I am trying to make a diagonal numpy array from:

[1,2,3,4,5,6,7,8,9]

Expected result:

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

What would be an efficient way of doing this?

MSeifert
  • 145,886
  • 38
  • 333
  • 352
delta27
  • 197
  • 1
  • 7

2 Answers2

2

You can use integer array indexing to set the specified elements of the output:

>>> import numpy as np
>>> a = [1,2,3,4,5,6,7,8,9]
>>> arr = np.zeros((9, 5), dtype=int)           # create empty array
>>> arr[np.arange(9), np.arange(2,11) % 5] = a  # insert a 
>>> arr
array([[0, 0, 1, 0, 0],
       [0, 0, 0, 2, 0],
       [0, 0, 0, 0, 3],
       [4, 0, 0, 0, 0],
       [0, 5, 0, 0, 0],
       [0, 0, 6, 0, 0],
       [0, 0, 0, 7, 0],
       [0, 0, 0, 0, 8],
       [9, 0, 0, 0, 0]])
MSeifert
  • 145,886
  • 38
  • 333
  • 352
  • `np.fill_diagonal` has a `wrap` mode like this, but it only works for the main diagonal. The other `diag` functions allow offsets, but no wrapping. – hpaulj May 11 '17 at 22:01
1

Inspired by np.fill_diagonal which can wrap, but not offset:

In [308]: arr=np.zeros((9,5),int)
In [309]: arr.flat[2:45:6]=np.arange(1,10)
In [310]: arr
Out[310]: 
array([[0, 0, 1, 0, 0],
       [0, 0, 0, 2, 0],
       [0, 0, 0, 0, 3],
       [0, 0, 0, 0, 0],
       [4, 0, 0, 0, 0],
       [0, 5, 0, 0, 0],
       [0, 0, 6, 0, 0],
       [0, 0, 0, 7, 0],
       [0, 0, 0, 0, 8]])

(though for some reason this has the 4th all zero row).

def fill_diagonal(a, val, wrap=False):
    ...       
    step = a.shape[1] + 1
    # Write the value out into the diagonal.
    a.flat[:end:step] = val
hpaulj
  • 221,503
  • 14
  • 230
  • 353
  • you lost the `9` and the fourth row seem wrong. Using a `step` works great until you have one value in the last column and expect the next value in the next row, first column :) – MSeifert May 11 '17 at 22:18
  • Yes, I"m puzzled about that. I'll leave the answer up in case someone sees the flaw in this 'strided' solution. – hpaulj May 11 '17 at 22:21
  • well, the flaw is that you need different "steps". 6 for row 1->2, 2->3, 4->5, 5->6, 6->7, 7->8 but only 1 for row 3->4 and 8->9. – MSeifert May 11 '17 at 22:23
  • This style of wrapping was the default for `fill_diagonal` until a version 1.6. – hpaulj May 11 '17 at 22:25
  • So when lined up just right the stride of 6 skips a whole row of length 5. – hpaulj May 11 '17 at 22:47