Add a column to an array with values from a position in another array if rows match-CodePudding

I have two arrays, one looks like this:

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

The other looks like this:

[[1 2 1 0 2 0]
[1 1 1 0 2 0]
[1 1 1 0 2 0]
[1 2 1 0 2 0]
[0 3 2 2 0 0]
[0 1 2 1 0 0]
[0 2 1 2 1 0]
[0 0 0 1 1 0]
[0 0 0 0 1 0]
[0 0 0 0 0 1]
     ...
 [0 3 2 2 0 0]
[0 1 2 1 0 0]
[0 2 1 2 1 0]
[0 0 0 1 1 0]
[0 0 0 0 1 0]
[0 0 0 0 0 1]]

Whenever a row in the second array matches the first six values in the first array I need to add the last element of the first array (the 7th element) at the end of the row of the second array that matches and when it doesn't match add a 0. The result would look like this:

[[1 2 1 0 2 0 1]
[1 1 1 0 2 0 0]
[1 1 1 0 2 0 0]
[1 2 1 0 2 0 1]
[0 3 2 2 0 0 0]
[0 1 2 1 0 0 2]
[0 2 1 2 1 0 0]
[0 0 0 1 1 0 4]
[0 0 0 0 1 0 5]
[0 0 0 0 0 1 6]
     ...
[0 3 2 2 0 0 0]
[0 1 2 1 0 0 2]
[0 2 1 2 1 0 0]
[0 0 0 1 1 0 4]
[0 0 0 0 1 0 5]
[0 0 0 0 0 1 6]]

CodePudding user response：

You could use:

import numpy as np
m = (B == A[:,None,:6]).all(2)
new_A = np.c_[B, np.where(m.any(0), np.take(A[:,6], m.argmax(0)), 0)]

How it works:

1- use broadcasting to compare B with all combinations of rows of A (limited to first 6 columns), and build a mask

2- Using numpy.where to check the condition: if at least 1 row in A matches, use numpy.argmax to get the index of the first match, and numpy.take to get the value from A's last column. Else, assign 0.

3- concatenate B and the newly build column

output:

array([[1, 2, 1, 0, 2, 0, 1],
       [1, 1, 1, 0, 2, 0, 0],
       [1, 1, 1, 0, 2, 0, 0],
       [1, 2, 1, 0, 2, 0, 1],
       [0, 3, 2, 2, 0, 0, 0],
       [0, 1, 2, 1, 0, 0, 2],
       [0, 2, 1, 2, 1, 0, 0],
       [0, 0, 0, 1, 1, 0, 4],
       [0, 0, 0, 0, 1, 0, 5],
       [0, 0, 0, 0, 0, 1, 6],
       [0, 3, 2, 2, 0, 0, 0],
       [0, 1, 2, 1, 0, 0, 2],
       [0, 2, 1, 2, 1, 0, 0],
       [0, 0, 0, 1, 1, 0, 4],
       [0, 0, 0, 0, 1, 0, 5],
       [0, 0, 0, 0, 0, 1, 6]])

inputs:

A = [[1, 2, 1, 0, 2, 0, 1],
     [1, 2, 1, 0, 2, 0, 1],
     [1, 2, 1, 0, 2, 0, 1],
     [1, 2, 1, 0, 2, 0, 1],
     [0, 1, 2, 1, 0, 0, 2],
     [0, 1, 2, 1, 0, 0, 2],
     [0, 0, 1, 0, 1, 0, 3],
     [0, 0, 0, 1, 1, 0, 4],
     [0, 0, 0, 0, 1, 0, 5],
     [0, 0, 0, 0, 0, 1, 6]]
A = np.array(A)
 
B = [[1, 2, 1, 0, 2, 0],
     [1, 1, 1, 0, 2, 0],
     [1, 1, 1, 0, 2, 0],
     [1, 2, 1, 0, 2, 0],
     [0, 3, 2, 2, 0, 0],
     [0, 1, 2, 1, 0, 0],
     [0, 2, 1, 2, 1, 0],
     [0, 0, 0, 1, 1, 0],
     [0, 0, 0, 0, 1, 0],
     [0, 0, 0, 0, 0, 1],
     [0, 3, 2, 2, 0, 0],
     [0, 1, 2, 1, 0, 0],
     [0, 2, 1, 2, 1, 0],
     [0, 0, 0, 1, 1, 0],
     [0, 0, 0, 0, 1, 0],
     [0, 0, 0, 0, 0, 1]]
B = np.array(B)