Find indexes in list based on category value-CodePudding

I am trying to find all the indexes categorized by the a given list of values.

For example:

n = [0,1,2,3]
b = [1,1,1,2,3,0,0,1,2]

I need a way to automate the extraction of the indexes in b that correspond to all the values in n, and output something like:

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

What is an approach I can do for this? Thanks

CodePudding user response：

One approach is to use a dictionary to with the keys as the elements of n and the indices as values, as:

n = [0, 1, 2, 3]
b = [1, 1, 1, 2, 3, 0, 0, 1, 2]

d = {ni: [] for ni in n}
for i, bi in enumerate(b):
    if bi in d:
        d[bi].append(i)

print(d)

Output

{0: [5, 6], 1: [0, 1, 2, 7], 2: [3, 8], 3: [4]}

CodePudding user response：

Use numpy.

import numpy as np

[np.where(np.array(b) == i)[0].tolist() for i in n]

Output:

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

CodePudding user response：

I would recommend using enumerate. This will let you compare each entry in b with the values from n, while also having access to the index where the entry in b lies. See the python code below:

n = [0,1,2,3]
b = [1,1,1,2,3,0,0,1,2]

for value in n:
  list = []
  for idx, entry in enumerate(b):
    if value == entry:
      list.append(idx)

  print(str(value)   ": "   str(list))

This prints the following to the console.

0: [5, 6]

1: [0, 1, 2, 7]

2: [3, 8]

3: [4]

CodePudding user response：

for y in n:
    print('{}:'.format(y),[i for i, x in enumerate(b) if x == y])

#output
0: [5, 6]
1: [0, 1, 2, 7]
2: [3, 8]
3: [4]