Add a column to a numpy array that counts if rows change-CodePudding

I have the following array:

[[1 2 1 0 2 0]
[1 2 1 0 2 0]
[1 2 1 0 2 0]
[1 2 1 0 2 0]
[0 1 2 1 0 0]
[0 1 2 1 0 0]
[0 0 1 0 1 0]
[0 0 0 1 1 0]
[0 0 0 0 1 0]
[0 0 0 0 0 1]]

I need to add a column to this array that adds a number whenever the values in the rows change starting with number 3. So the result would look like this:

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

Thank you

CodePudding user response：

If a is your array as:

a = np.array([[1, 2, 1, 0, 2, 0], [1, 2, 1, 0, 2, 0], [1, 2, 1, 0, 2, 0], [1, 2, 1, 0, 2, 0],
              [0, 1, 2, 1, 0, 0], [0, 1, 2, 1, 0, 0], [0, 0, 1, 0, 1, 0], [0, 0, 0, 1, 1, 0],
              [0, 0, 0, 0, 1, 0], [0, 0, 0, 0, 0, 1]])

using the following code will get you the results:

n = 3
a = a.tolist()
for i, j in enumerate(a):
    if i == 0:
        j.append(n)
    elif i > 0 and j == a[i-1][:-1]:
        j.append(n)
    else:
        n  = 1
        j.append(n)

# a = np.array(a)

which will give:

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