joining two numpy arrays by the elements of the first array

Question

a = np.array([[383, 383, 384, 384, 384, 384, 384, 384, 384, 421, 422, 422, 422, 422, 423, 423, 423, 423, 423, 423, 479, 480, 481, 482, 483, 485, 485, 485, 485, 485, 485, 485, 513, 515, 517]])
b = np.array([384, 423, 485])
np.where(a == b[:,None])

returns me two arrays

[0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2]

[2, 3, 4, 5, 6, 7, 8, 14, 15, 16, 17, 18, 19, 25, 26, 27, 28, 29, 30, 31]

and i want to join?(i don't know how to call that operation) the two lists by the values of the first one , so the output will be like

[[2, 3, 4, 5, 6, 7, 8], [14, 15, 16, 17, 18, 19], [25, 26, 27, 28, 29, 30, 31]]

how can i do it?

(my final mission is to get the central value of each list in the last array)

I think this would likely be called a "grouping" operation more than a joining one. — Henry Ecker, Jul 23 '21 at 23:46
In case the duplicate is unclear `c, d = np.where(a == b[:, None])` then `np.split(d, np.unique(c, return_index=True)[1][1:])` — Henry Ecker, Jul 24 '21 at 00:03

score 0 · Accepted Answer · answered Jul 24 '21 at 00:06

You can do it like this:

>>> [np.where(a==x)[1] for x in b]
[array([2, 3, 4, 5, 6, 7, 8], dtype=int64),
 array([14, 15, 16, 17, 18, 19], dtype=int64),
 array([25, 26, 27, 28, 29, 30, 31], dtype=int64)]

If you want the output as lists instead of arrays:

>>> [list(np.where(a==x)[1]) for x in b]
[[2, 3, 4, 5, 6, 7, 8], [14, 15, 16, 17, 18, 19], [25, 26, 27, 28, 29, 30, 31]]

joining two numpy arrays by the elements of the first array

1 Answers1