Both increment and decrement in While loop in matrix in python

Question

I created a 15 x 15 matrix in a pandas dataframe type. I want to make changes in some cells, following this logic:

• The diagonal of the matrix is set to 0
• In each row, if 0 appears in any position, the values in the next/ previous 5 columns should be updated to 1. (df.iloc[i][j] = 0 --> df.iloc[i][j+5] = 1 or df.iloc[i][j-5] = 1

Input (every cell is 0.9):

row, col = 15, 16
a = pd.DataFrame.from_records([[0.9]*col]*row)
a = a.loc[ : , a.columns != 0]

Expected output:

Below is my current script (which keeps running without an end):

row, col = 15, 16
a = pd.DataFrame.from_records([[0.9]*col]*row)
a = a.loc[ : , a.columns != 0]
column3 = list(a)
for i, j in a.iterrows():
    for k in column3:
        if k == i + 1:
            a.iloc[i][k] = 0


for i, j in a.iterrows():
    for k in column3:
        step = +5 if k <=5 else -5
        while k <= len(a.index):
            if a.iloc[i][k] == 0:
                k += step
                a.iloc[i][k] = 1

pd.set_option('display.max_rows', None)
pd.set_option('display.max_columns', None)
a 

Many thanks in advance for your help!

Answer 1

I would use numpy here and build a mask with roll:

# distance to 0
N = 5

# replace diagonal with 0
np.fill_diagonal(a.values, 0)

# build mask
m = (a==0).to_numpy()

# apply mask iteratively
for i in range(1, a.shape[1]//N):
    a[np.roll(m, i*N, axis=1)] = 1

Variant using pandas' shift:

N = 5

np.fill_diagonal(a.values, 0)
m = (a==0)

for i in range(1, a.shape[1]//N):
    a[m.shift(i*N, axis=1, fill_value=False)] = 1
    a[m.shift(-i*N, axis=1, fill_value=False)] = 1

Output:

     1    2    3    4    5    6    7    8    9    10   11   12   13   14   15
0   0.0  0.9  0.9  0.9  0.9  1.0  0.9  0.9  0.9  0.9  1.0  0.9  0.9  0.9  0.9
1   0.9  0.0  0.9  0.9  0.9  0.9  1.0  0.9  0.9  0.9  0.9  1.0  0.9  0.9  0.9
2   0.9  0.9  0.0  0.9  0.9  0.9  0.9  1.0  0.9  0.9  0.9  0.9  1.0  0.9  0.9
3   0.9  0.9  0.9  0.0  0.9  0.9  0.9  0.9  1.0  0.9  0.9  0.9  0.9  1.0  0.9
4   0.9  0.9  0.9  0.9  0.0  0.9  0.9  0.9  0.9  1.0  0.9  0.9  0.9  0.9  1.0
5   1.0  0.9  0.9  0.9  0.9  0.0  0.9  0.9  0.9  0.9  1.0  0.9  0.9  0.9  0.9
6   0.9  1.0  0.9  0.9  0.9  0.9  0.0  0.9  0.9  0.9  0.9  1.0  0.9  0.9  0.9
7   0.9  0.9  1.0  0.9  0.9  0.9  0.9  0.0  0.9  0.9  0.9  0.9  1.0  0.9  0.9
8   0.9  0.9  0.9  1.0  0.9  0.9  0.9  0.9  0.0  0.9  0.9  0.9  0.9  1.0  0.9
9   0.9  0.9  0.9  0.9  1.0  0.9  0.9  0.9  0.9  0.0  0.9  0.9  0.9  0.9  1.0
10  1.0  0.9  0.9  0.9  0.9  1.0  0.9  0.9  0.9  0.9  0.0  0.9  0.9  0.9  0.9
11  0.9  1.0  0.9  0.9  0.9  0.9  1.0  0.9  0.9  0.9  0.9  0.0  0.9  0.9  0.9
12  0.9  0.9  1.0  0.9  0.9  0.9  0.9  1.0  0.9  0.9  0.9  0.9  0.0  0.9  0.9
13  0.9  0.9  0.9  1.0  0.9  0.9  0.9  0.9  1.0  0.9  0.9  0.9  0.9  0.0  0.9
14  0.9  0.9  0.9  0.9  1.0  0.9  0.9  0.9  0.9  1.0  0.9  0.9  0.9  0.9  0.0

Output as image for clarity: