CODEWITHSUNDEEP
pythonpandasdata-analysis
Neel
Neel

Reputation: 10143

Pandas DataFrame as lookup table with embedded lists

I have a dataframe with the following structure:

     A          B
    [1, 2, 3] [a, b, c]
    [4, 5, 6] [d, e, f]

I want to query the dataframe such that when I input 1, it should return [a,b,c]. Similarly, querying for 6 should return [d, e, f]. What's the most readable way to do it?

Upvotes: 2

Views: 100

Answers (2)

Andy L.
Andy L.

Reputation: 25239

Use map and loc

n = 1
df.loc[df.A.map(lambda x: n in x), 'B']

Out[209]:
0    [a, b, c]
Name: B, dtype: object

Upvotes: 2

Paul H
Paul H

Reputation: 68146

I would explode your dataframe:

df = df.explode('A').set_index('A')

n = 1
print(df.loc[n, 'B'])

prints:

['a', 'b', 'c']

Upvotes: 0

Related Questions