Nakeuh
Nakeuh

Reputation: 1909

Retrieve array elements from a list of index

I have an array my_array containing some tuple data.

I have another array my_array_values of same length containing some integer values.

For each unique value of my_array_values, I want to retrieve a list of values from my_array, that are in the same index as this value in my_array_values. Here is my code and the expected behavior :

my_array = np.array([('AA','11'),('BB','22'),('CC','33'),('DD','44'),('EE','55'),('FF','66')])
my_array_values = np.array([1,2,3,1,3,2])
my_array_values_unique = np.array([1,2,3])

for v in my_array_values_unique:
    print(np.take(my_array, np.where(my_array_values == v)))

Expected behavior :

[('AA', '11'), ('DD', '44')]
[('BB', '22'), ('FF', '66')]
[('CC', '33'), ('EE', '55')]

But actually, my code gives me the following output:

[['AA' '22']]
[['11' '33']]
[['BB' 'CC']]

Can someone explain to me how I can get the correct output?

Upvotes: 1

Views: 79

Answers (3)

Tomerikoo
Tomerikoo

Reputation: 19405

You don't need to use take or where at all. Equality check on an array returns a boolean array which is a valid indexing array:

for v in my_array_values_unique:
    print(my_array[my_array_values == v])

And this prints:

[['AA' '11']
 ['DD' '44']]
[['BB' '22']
 ['FF' '66']]
[['CC' '33']
 ['EE' '55']]

If numpy is not specificly required, this can be easily done using lists as well:

lst = [('AA', '11'), ('BB', '22'), ('CC', '33'), ('DD', '44'), ('EE', '55'), ('FF', '66')]
idxs = [1, 2, 3, 1, 3, 2]

for v in set(idxs):
    print([tup for idx, tup in zip(idxs, lst) if idx == v])

Gives:

[('AA', '11'), ('DD', '44')]
[('BB', '22'), ('FF', '66')]
[('CC', '33'), ('EE', '55')]

Another, more efficient, way would be to use defaultdict in order to loop the list once, instead of once for every unique value:

import collections

mapping = collections.defaultdict(list)
for tup, idx in zip(lst, idxs):
    mapping[idx].append(tup)

for lst in mapping.values():
    print(lst)
  • gives the same result as before

Upvotes: 2

Junior Yao
Junior Yao

Reputation: 31

I am not sure but, If you are using NumPy array, it will return a list of an array within an array, so you won't get any tuple from it, until and unless you manipulate it again

Upvotes: 0

Sanchit.Jain
Sanchit.Jain

Reputation: 608

Please use axis=0 when using np.take in this case. By default, it flattens out your array that's why you are getting 'AA' and '22' for case '1'.

Upvotes: 2

Related Questions