Pandas Dataframe - How to stack two column values into one list?

Question

I have a dataframe with three columns (from, to, w):

      from to w
0     0   1  0.670820
1     0   5  0.612372
2     0   2  0.612372
3     0   3  0.577350
4     0   4  0.408248

How can I get the values in a list from the two columns (from, to) based on a given id? For example, in the above example, the list will be [1,5,2,3,4] for a given id 0.

Note that id may appear in from or two columns. In the example, below the expected list will be [2,4,0,3,5] if the given id is 1.

     from  to   w
0     1   2  0.730297
1     1   4  0.730297
2     0   1  0.670820
3     1   3  0.516398
4     1   5  0.365148

I iterate the rows in the dataframe to generate the list:

myarr =[]
    for index, row in dftemp1.iterrows():
        from_id = row['from']
        to_id = row['to']
        if (from_id!=target_id):
            myarr.append(from_id)
        if (to_id!=target_id):
            myarr.append(to_id)

I wonder if there is a simpler way to achieve the results. Any help would be appreciated.

score 1 · Accepted Answer · answered Jan 10 '18 at 08:18

Using your second example -

df

   from  to         w
0     1   2  0.730297
1     1   4  0.730297
2     0   1  0.670820
3     1   3  0.516398
4     1   5  0.365148

You can argsort the values in the first two columns depending on whether the value is equal to your ID or not.

v = df.iloc[:, :-1].values

i = np.arange(len(df))[:, None]
j = np.argsort(v == 1, axis=1)   # replace `1` with your ID

v[i, j][:, 0].tolist()
[2, 4, 0, 3, 5]

Pandas Dataframe - How to stack two column values into one list?

1 Answers1