In [1]:

    

import numpy as np


    

In [2]:

    

a = np.arange(9).reshape((3, 3))


    

In [3]:

    

print(a)


    

    




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

In [4]:

    

print(np.diag(a))


    

    




[0 4 8]

In [5]:

    

print(np.diag(a, k=1))


    

    




[1 5]

In [6]:

    

print(np.diag(a, k=2))


    

    




[2]

In [7]:

    

print(np.diag(a, k=3))


    

    




[]

In [8]:

    

print(np.diag(a, k=-1))


    

    




[3 7]

In [9]:

    

print(np.diag(a, k=-2))


    

    




[6]

In [10]:

    

print(np.diag(a, k=-3))


    

    




[]

In [11]:

    

a = np.arange(12).reshape((3, 4))


    

In [12]:

    

print(a)


    

    




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

In [13]:

    

print(np.diag(a))


    

    




[ 0  5 10]

In [14]:

    

print(np.diag(a, k=1))


    

    




[ 1  6 11]

In [15]:

    

print(np.diag(a, k=-1))


    

    




[4 9]

In [16]:

    

a = np.arange(27).reshape((3, 3, 3))


    

In [17]:

    

print(a)


    

    




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

 [[ 9 10 11]
  [12 13 14]
  [15 16 17]]

 [[18 19 20]
  [21 22 23]
  [24 25 26]]]

In [18]:

    

# print(np.diag(a))
# ValueError: Input must be 1- or 2-d.


    

In [19]:

    

a = np.arange(9).reshape((3, 3))


    

In [20]:

    

print(a)


    

    




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

In [21]:

    

a_diag = np.diag(a)


    

In [22]:

    

print(a_diag)


    

    




[0 4 8]

In [23]:

    

# a_diag[0] = 100
# ValueError: assignment destination is read-only


    

In [24]:

    

a_diag.flags.writeable = True


    

In [25]:

    

a_diag[0] = 100


    

In [26]:

    

print(a_diag)


    

    




[100   4   8]

In [27]:

    

print(a)


    

    




[[100   1   2]
 [  3   4   5]
 [  6   7   8]]

In [28]:

    

a_diag_copy = np.diag(a).copy()


    

In [29]:

    

print(a_diag_copy)


    

    




[100   4   8]

In [30]:

    

a_diag_copy[1] = 100


    

In [31]:

    

print(a_diag_copy)


    

    




[100 100   8]

In [32]:

    

print(a)


    

    




[[100   1   2]
 [  3   4   5]
 [  6   7   8]]