Efficient/Pythonic way to create lists from pandas Dataframe column

Question

I've a dataframe as below.

df = pd.DataFrame({
    'code' : [1,   2,  3,  4,  5,  6,  7,  8,  9,  10],
    'Tag' :  ['A','B','C','D','B','C','D','A','D','C']
})

+------+-----+
| code | Tag |
+------+-----+
|   1  |  A  |
+------+-----+
|   2  |  B  |
+------+-----+
|   3  |  C  |
+------+-----+
|   4  |  D  |
+------+-----+
|   5  |  B  |
+------+-----+
|   6  |  C  |
+------+-----+
|   7  |  D  |
+------+-----+
|   8  |  A  |
+------+-----+
|   9  |  D  |
+------+-----+
|  10  |  C  |
+------+-----+

My objective is to create code lists based on the common items in the Tag column as below.

codes_A = [1,8]
codes_B = [2,5]
codes_C = [3,6,10]
codes_D = [4,7,9]

How I'm doing it right now is

codes_A = df[df['Tag'] == 'A']['code'].to_list()
codes_B = df[df['Tag'] == 'B']['code'].to_list()
codes_C = df[df['Tag'] == 'C']['code'].to_list()
codes_D = df[df['Tag'] == 'D']['code'].to_list()

This code does the job. But, as you can see this is very cumbersome and inefficient. I'm repeating the same code multiple times and also repeating when I want to create new lists.

is there a more efficient and pythonic way to do this in pandas or numpy?

jezrael · Accepted Answer · 2021-03-04 09:44:18Z

3

Create dictionary of list, becasue variable names are not recommended:

d = df.groupby('Tag')['code'].agg(list).to_dict()
print (d)
{'A': [1, 8], 'B': [2, 5], 'C': [3, 6, 10], 'D': [4, 7, 9]}

Then for list lookup by keys in dict, but no assign to variable name:

print (d['A'])
[1, 8]

So practically it means in your code if use codes_A then it change to d['A'], similar for all variables.

But if really need it:

for k, v in d.items():
    globals()[f'code_{k}'] = v
    
print (code_A)
[1, 8]

edited Mar 4, 2021 at 9:44

answered Mar 4, 2021 at 9:25

jezrael

868k103 gold badges1.4k silver badges1.3k bronze badges

Sign up to request clarification or add additional context in comments.

3 Comments

Tommy Over a year ago

I want to create separate lists to be used in my other codes

jezrael Over a year ago

@Tommy - hmmm, it is posible by globals, but it is not pythonic way.

Dorian Turba Over a year ago

globals is a bad thing in a lot of language. A good explanation: youtu.be/UEuXQjPUwcw

Collectives™ on Stack Overflow

Efficient/Pythonic way to create lists from pandas Dataframe column

1 Answer 1

3 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

3 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related