Subset according to their label

user_14335678 0
Viewed 62

Suppose I have two lists: letters = [a, b, c, d, e, f, g, h, i] and digits = [0, 1, 1, 2, 1, 0, 2, 2, 1, 0]. I would like the final result to be output = [[a, f, i], [b, c, e, i], [d, g, h]].

Here 0, 1 and 2 are just different classes. For instance,a is from class 0 while b and c are from class 1. I just need to put the letters in a sublist according to their class.

I think I can you zip() here and a list comprehension, but I am not sure how to do that. How can I do that using numpy?

zip list-comprehension numpy python
5 Answers

You can use zip() and temporary dictionary:

letters = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i']
digits = [0, 1, 1, 2, 1, 0, 2, 2, 1, 0]

tmp = {}
for d, l in zip(digits, letters):
    tmp.setdefault(d, []).append(l)

out = []
for k in sorted(tmp):
    out.append(tmp[k])

print(out)

Prints:

[['a', 'f'], ['b', 'c', 'e', 'i'], ['d', 'g', 'h']]

Or: Another version (using itertools.groupby):

from itertools import groupby

out = []
for _, g in groupby(sorted(zip(digits, letters)), lambda k: k[0]):
    out.append([l for _, l in g])

print(out)
user_100359850

numpy solution is also possible:

digits = np.array([0, 1, 1, 2, 1, 0, 2, 2, 0])
letters = np.array(['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i'])
argidx = np.argsort(digits)
digits, letters = digits[argidx], letters[argidx]
markers = np.diff(digits, prepend=digits[0])
marker_idx, = np.nonzero(markers)
np.split(letters, marker_idx)

Output:

[array(['a', 'f', 'i'], dtype='<U1'), array(['b', 'c', 'e'], dtype='<U1'), array(['d', 'g', 'h'], dtype='<U1')]
user_30448250

Since you tagged numpy (I personally prefer pandas groupby since it is clean and meant for this purpose):

d,l = list(zip(*(sorted(zip(digits,letters)))))
d = np.array(d)
idx = np.flatnonzero(np.r_[True, d[:-1] != d[1:], True])
output = [list(l[i:j]) for i,j in zip(idx[:-1],idx[1:])]
#[['a', 'f'], ['b', 'c', 'e', 'i'], ['d', 'g', 'h']]
user_49759810

Another variation on the answer (here with groupby and itemgetter):

from itertools import groupby
from operator import itemgetter

letters = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i']
digits = [0, 1, 1, 2, 1, 0, 2, 2, 1, 0]

combined_letters_digits = sorted(zip(letters, digits), key=itemgetter(1))
letter_groups = groupby(combined_letters_digits, itemgetter(1))
out = [[item[0] for item in group_data] for (key, group_data) in letter_groups]

print(out)
user_15708540

much more concise, no loops with pandas

import pandas as pd

s = pd.Series(
    ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i'],
    index=[0, 1, 1, 2, 1, 0, 2, 2, 1],
    name='letters'
)

[*s.groupby(level=0).agg(list)]


[['a', 'f'], ['b', 'c', 'e', 'i'], ['d', 'g', 'h']]
user_104030900
Related