How To Iterate Over All Elements Of A 2D Matrix Using Only One Loop Using Python

I know you can iterate over a 2d matrix using two indexes like this: import numpy as np A = np.zeros((10,10)) for i in range(0,10): for j in range(0,10): if (i==j):

Solution 1:

In [129]: A = np.zeros((10,10), int)
     ...: for i in range(0,10):
     ...:     for j in range(0,10):
     ...:         if (i==j):
     ...:             A[i,j] = 1
     ...:         if (i+1 ==j):
     ...:             A[i,j] = 2
     ...:         if (i-1==j):
     ...:             A[i,j] = 3
     ...: 

You should have shown the resulting A:

In [130]: A
Out[130]: 
array([[1, 2, 0, 0, 0, 0, 0, 0, 0, 0],
       [3, 1, 2, 0, 0, 0, 0, 0, 0, 0],
       [0, 3, 1, 2, 0, 0, 0, 0, 0, 0],
       [0, 0, 3, 1, 2, 0, 0, 0, 0, 0],
       [0, 0, 0, 3, 1, 2, 0, 0, 0, 0],
       [0, 0, 0, 0, 3, 1, 2, 0, 0, 0],
       [0, 0, 0, 0, 0, 3, 1, 2, 0, 0],
       [0, 0, 0, 0, 0, 0, 3, 1, 2, 0],
       [0, 0, 0, 0, 0, 0, 0, 3, 1, 2],
       [0, 0, 0, 0, 0, 0, 0, 0, 3, 1]])

So you have set 3 diagonals:

In [131]: A[np.arange(10),np.arange(10)]
Out[131]: array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1])
In [132]: A[np.arange(9),np.arange(1,10)]
Out[132]: array([2, 2, 2, 2, 2, 2, 2, 2, 2])
In [133]: A[np.arange(1,10),np.arange(9)]
Out[133]: array([3, 3, 3, 3, 3, 3, 3, 3, 3])

The key to eliminating loops in numpy is to get a big picture of the task, rather than focusing on the iterative steps.

There are various tools for making a diagonal array. One is np.diag, which can be used thus:

In [139]: np.diag(np.ones(10,int),0)+
          np.diag(np.ones(9,int)*2,1)+
          np.diag(np.ones(9,int)*3,-1)
Out[139]: 
array([[1, 2, 0, 0, 0, 0, 0, 0, 0, 0],
       [3, 1, 2, 0, 0, 0, 0, 0, 0, 0],
       [0, 3, 1, 2, 0, 0, 0, 0, 0, 0],
       [0, 0, 3, 1, 2, 0, 0, 0, 0, 0],
       [0, 0, 0, 3, 1, 2, 0, 0, 0, 0],
       [0, 0, 0, 0, 3, 1, 2, 0, 0, 0],
       [0, 0, 0, 0, 0, 3, 1, 2, 0, 0],
       [0, 0, 0, 0, 0, 0, 3, 1, 2, 0],
       [0, 0, 0, 0, 0, 0, 0, 3, 1, 2],
       [0, 0, 0, 0, 0, 0, 0, 0, 3, 1]])

Or adapting [131] etc

In [140]: A = np.zeros((10,10), int)
     ...: A[np.arange(10),np.arange(10)]=1
     ...: A[np.arange(9),np.arange(1,10)]=2
     ...: A[np.arange(1,10),np.arange(9)]=3

Solution 2:

Because your only executing code when i == j, you can just use:

for i in range(0,10):
    A[i,i] = 4

Solution 3:

You can always collapse multiple loops into one by calculating the components each iteration with the modulo operator like so:

import numpy as np

A = np.zeros((10,10))

for x in range(100):
    i = math.floor(x/10)
    j = x % 10
    if (i==j):
        A[i,j] = 1
    if (i+1 ==j):
        A[i,j] = 2
    if (i-1==j):
        A[i,j] = 3

With only i==j it could be even simpler:

for i in range(10):
    A[i,i] = 4