Need a list of row-values in pandas

user_8854526 0
Viewed 129

What I have, and what I need

I have a pandas DataFrame p with cols 'a', 'b', 'c' (col names stored in pc).

From that I would like to create a DataFrame pn of the same shape, but each cell as a list of values from selected rows. The DataFrame n tells me which rows to select from p for each row in pn.

import pandas as pd

pc = ['a', 'b', 'c']
p  = pd.DataFrame([[11, 12, 13],
                   [21, 22, 23]],
                  columns=pc,
                  index=[1001,
                         1002])

n  = pd.DataFrame([[[1001]      ],
                   [[1001, 1002]]],
                  columns=['sel_row'],
                  index=[1001,
                         1002])

What I could (and want to) achieve

The farthest I could get... gives me a list of cols, rather than rows. So, am I confusing the nested for loops ?

pn = pd.DataFrame([ [p.loc[ix, pc].values for ix in n.loc[indx].values[0]] 
                    for indx in n.index ])

print (pn)

# The actual output:
#               0             1
# 0  [11, 12, 13]          None
# 1  [11, 12, 13]  [21, 22, 23]

# The required output:
#           0         1         2
# 0  [11]      [12]      [13]
# 1  [11, 21]  [12, 22]  [13, 23]

Stray thoughts

Maybe I should also iterate something like p.loc[ix, c] ... for c in pc... but how can there be 3 loops ??

A further (optional) wish

Is this possible with lambda too ? My intuition is: that would be faster-- but not sure !

Thanks for going through the question or any help offered.

pandas dataframe list python
3 Answers

You can explode the n, use that to slice p and groupby:

s = n['sel_row'].explode()
p.loc[s].groupby(s.index).agg(list)

Output:

             a         b         c
1001      [11]      [12]      [13]
1002  [11, 21]  [12, 22]  [13, 23]
user_42384080

IIUC, you could do:

data = [[[*x] for x in zip(*p.loc[idxs].values)] for idxs in n['sel_row']]

result = pd.DataFrame(data=data, columns=p.columns, index=p.index)
print(result)

Output

             a         b         c
1001      [11]      [12]      [13]
1002  [11, 21]  [12, 22]  [13, 23]
user_40015920

You can write a custom function here.

pc = ['a', 'b', 'c']
p  = pd.DataFrame([[11, 12, 13],
                   [21, 22, 23]],
                  columns=pc,
                  index=[1001,
                         1002])

n  = pd.DataFrame([[[1001]      ],
                   [[1001, 1002]]],
                  columns=['sel_row'],
                  index=[1001,
                         1002])
def f(idx):
    return pd.Series(p.loc[idx, :].values.T.tolist())

n.sel_row.apply(f)

             0         1         2
1001      [11]      [12]      [13]
1002  [11, 21]  [12, 22]  [13, 23]

With lambda could rewrite above as:

n.sel_row.apply(lambda idx: pd.Series(p.loc[idx, :].values.T.tolist()))
user_124164530
Related