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

Question

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

mozway · Accepted Answer · 2021-11-14T06:53:03.207

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)

Thank you so much for your answer, this was the last piece of the puzzle. However, could you assign the final array to a variable like this: r = np.c_[B, np.where(m.any(0), np.take(A[:,6], m.argmax(0)), 0)] This is because I wasn't sure if the output was a new array or a modified in-place version of m. — Gus, Nov 13 '21 at 18:29
Now I know that the output is a new array but I think that your answer can be more clear if you give this array a letter. This is for some other people like myself that yet don't know if your command is a new array or a method that modifies m. In this last scenario, the array in the output is m and not a new one, however, since it is a new array to avoid confusion you can give a new letter and everyone will know is not m but a completely new array. When I first saw your answer I thought the output array was m so when I printed m the result didn't match, then I realize I need to assign to "r" — Gus, Nov 14 '21 at 01:05

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

1 Answers1