python - 获取影响原始数组的 2D numpy 数组的对角线

Question

我有一个二维 numpy 数组。我想通过创建其对角线元素数组并修改对角线数组来修改它，以便这些更改反射(reflect)在原始二维数组上。

我试过:

>>> a = np.ones(shape=(3,3))
>>> d1 = a[np.diag_indices_from(a)]
>>> d1
array([1., 1., 1.])
>>> d1[0] = 2
>>> d1
array([2., 1., 1.])
>>>a
array([[1., 1., 1.],
       [1., 1., 1.],
       [1., 1., 1.]])

可以看出，改变并没有影响到原来的数组。

有什么方法可以创建一个对角线阵列，它也会影响原始二维阵列吗？

编辑: 当我处理行或列时，我得到了我正在寻找的效果:

>>> row0 = a[0]
>>> row0[0] = 0
>>> a
array([[0., 1., 1.],
       [1., 1., 1.],
       [1., 1., 1.]])
>>> column0=a[:,0]
>>> column0[2]=3
>>> a
array([[0., 1., 1.],
       [1., 1., 1.],
       [3., 1., 1.]])

Answer 1

您可以使用 np.einsum 来精确获得您想要的:

.. versionadded:: 1.10.0

Views returned from einsum are now writeable whenever the input array
is writeable. For example, ``np.einsum('ijk...->kji...', a)`` will now
have the same effect as :py:func:`np.swapaxes(a, 0, 2) <numpy.swapaxes>`
and ``np.einsum('ii->i', a)`` will return a writeable view of the diagonal
of a 2D array.

a = np.ones((3,3))
b = np.einsum("ii->i",a)
b[:] = 2,3,4
a
# array([[2., 1., 1.],
#        [1., 3., 1.],
#        [1., 1., 4.]])