Создать двухрядную диагональную матрицу в numpy

Question

Я пытаюсь сделать специальную диагональную матрицу, которая выглядит следующим образом:

[[1,1,0,0,0,0],
 [0,0,1,1,0,0],
 [0,0,0,0,1,1]]

Это немного отличается от вопроса здесь: Сделать специальную диагональную матрицу в Numpy

Я пытался настроить решение, но не смог его получить.Цените любые советы о том, как этого добиться эффективно.

Answer 1

Не так элегантно, как в комментариях, но:

a=4 # number of rows
b=a*2 #number of columns
np.array((([1]*2+[0]*b)*a)[:-b]).reshape(a,b)
array([[1, 1, 0, 0, 0, 0, 0, 0],
       [0, 0, 1, 1, 0, 0, 0, 0],
       [0, 0, 0, 0, 1, 1, 0, 0],
       [0, 0, 0, 0, 0, 0, 1, 1]])

работает для любых a.