0

I have a dataframe and I would like to extract data by some conditions.
My dataframe is like this:

 A  B   D
e1  r2  a
e8  r7  a
e2  r2  a..b
e5  e10 c
e8  e12 c..a

For each value in the column D (a, b, c) I would like to get values from A and B like this:

a : [[e1, r2], [e8, r7], [e2,r2], [e8, e12]]
b : [[e2, r2]]
c : [[e5, e10], [e8, e12]]
....

This is what I've tried but I don't know how to extract the columns A and B.

df
l = ['a','b','c']
list_elements = {}
for i in l:
    liste_e = []
        for e in df['d']:
        if i.upper() in e:
            liste_e.append([e1, r2])
        ## extract a : [[e1, r2], [e8, r7], [e2,r2], [e8, e12]]

        list_elements[i] = liste_e
Georgy
  • 12,464
  • 7
  • 65
  • 73
noaai
  • 59
  • 1

2 Answers2

1

You can also create a dict with empty lists as values and append to it

l = ['a','b','c']
dic = dict([(key, []) for key in l])
for i in l:
    value = df.D.str.contains(i)
    dic[i].append(df.loc[value, ['A', 'B']].values)
Sumanth
  • 497
  • 4
  • 12
0

You can use .str.contains to get boolean masks for each predefined letter, get corresponding parts of A and B column using .loc, and convert to list of lists those parts of the dataframe that we obtained in the previous step.

Your dataframe:

>>> import pandas as pd
>>> df = pd.DataFrame(dict(A=['e1', 'e8', 'e2', 'e5', 'e8'],
                           B=['r2', 'r7', 'r2', 'e10', 'e12'],
                           D=['a', 'a', 'a..b', 'c', 'c..a']))
>>> df
     A  B   D
0   e1  r2  a
1   e8  r7  a
2   e2  r2  a..b
3   e5  e10 c
4   e8  e12 c..a

then

>>> letters = ['a', 'b', 'c']
>>> lists = {}
>>> for letter in letters:
        mask = df['D'].str.contains(letter)
        lists[letter] = df.loc[mask, ['A', 'B']].values.tolist()
>>> lists
{'a': [['e1', 'r2'], ['e8', 'r7'], ['e2', 'r2'], ['e8', 'e12']],
 'b': [['e2', 'r2']],
 'c': [['e5', 'e10'], ['e8', 'e12']]}
Georgy
  • 12,464
  • 7
  • 65
  • 73