根据第一个索引遍历数组

Question

I have two arrays and I am wanting to loop through a second array to only return arrays whose first element is equal to an element from another array.

 a = [10, 11, 12, 13, 14]
 b = [[9, 23, 45, 67, 56, 23, 54], [10, 8, 52, 30, 15, 47, 109], [11, 81, 
 152, 54, 112, 78, 167], [13, 82, 84, 63, 24, 26, 78], [18, 182, 25, 63, 96, 
 104, 74]]

I have two different arrays, a and b. I would like to find a way to look through each of the sub-arrays(?) within b in which the first value is equal to the values in array a to create a new array, c.

The result I am looking for is:

  c = [[10, 8, 52, 30, 15, 47, 109],[11, 81, 152, 54, 112, 78, 167],[13, 82, 84, 63, 24, 26, 78]]

Does Python have a tool to do this in a way Excel has MATCH()?

I tried looping in a manner such as:

 for i in a:
      if i in b:
          print (b)

But because there are other elements within the array, this way is not working. Any help would be greatly appreciated.

Further explanation of the problem:

a = [5, 6, 7, 9, 12]

I read in a excel file using XLRD (b_csv_data):

 Start  Count   Error   Constant    Result1 Result2 Result3 Result4
 5       41       0       45             23      54      66       19
 5.4     44       1       21             52      35       6       50
 6       16       1       42             95      39       1       13
 6.9     50       1       22             71      86      59       97
 7       38       1       43             50      47      83       67
  8      26       1       29             100     63      15       40
 9       46       0       28             85       9      27       81
 12      43       0       21             74      78      20       85

Next, I created a look to read in a select number of rows. For simplicity, this file above only has a few rows. My current file has about 100 rows.

for r in range (1, 7): #skipping headers and only wanting first few rows to start

     b_raw = b_csv_data.row_values(r) 
     b = np.array(b_raw) # I created this b numpy array from the line of code above

Answer 1

Use np.isin -

In [8]: b[np.isin(b[:,0],a)]
Out[8]: 
array([[ 10,   8,  52,  30,  15],
       [ 11,  81, 152,  54, 112],
       [ 13,  82,  84,  63,  24]])

With sorted a , we can also use np.searchsorted -

idx = np.searchsorted(a,b[:,0])
idx[idx==len(a)] = 0
out = b[a[idx] == b[:,0]]

If you have an array with different number of elements per row, which is essentially array of lists, you need to modify the slicing part. So, in that case, get the first off elements -

b0 = [bi[0] for bi in b]

Then, use b0 to replace all instances of b[:,0] in earlier posted methods.

Answer 2

Use list comprehension:

c = [l for l in b if l[0] in a]

Output:

[[10, 8, 52, 30, 15], [11, 81, 152, 54, 112], [13, 82, 84, 63, 24]]

If your list or array s are considerably large, using numpy.isin can be significantly faster:

b[np.isin(b[:, 0], a), :]

Benchmark:

a = [10, 11, 12, 13, 14]
b = [[9, 23, 45, 67, 56], [10, 8, 52, 30, 15], [11, 81, 152, 54, 112], 
 [13, 82, 84, 63, 24], [18, 182, 25, 63, 96]]

list_comp, np_isin = [], []
for i in range(1,100):
    a_test = a * i
    b_test = b * i
    list_comp.append(timeit.timeit('[l for l in b_test if l[0] in a_test]', number=10, globals=globals()))
    a_arr = np.array(a_test)
    b_arr = np.array(b_test)
    np_isin.append(timeit.timeit('b_arr[np.isin(b_arr[:, 0], a_arr), :]', number=10, globals=globals()))

While it is not clear and concise, I would recommend using list comprehension if the b is shorter than 100. Otherwise, numpy is your way to go.

Answer 3

You are doing it reverse. It is better to loop through the elements of b array and check if it is present in a. If yes then print that element of b. See the answer below.

a = [10, 11, 12, 13, 14]
b = [[9, 23, 45, 67, 56, 23, 54], [10, 8, 52, 30, 15, 47, 109], [11, 81, 152, 54, 112, 78, 167], [13, 82, 84, 63, 24, 26, 78], [18, 182, 25, 63, 96, 104, 74]]

for bb in b:  # if you want to check only the first element of b is in a
    if bb[0] in a:
            print(bb)

for bb in b:   # if you want to check if any element of b is in a
    for bbb in bb:
        if bbb in a:
            print(bb)

Output:

[10, 8, 52, 30, 15, 47, 109]
[11, 81, 152, 54, 112, 78, 167]
[13, 82, 84, 63, 24, 26, 78]