Search code examples
arraysnumpymatrixcumsum

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


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


Solution

  • 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]]