split based on comma and create new data frame in Python

Viewed 49

Let's say I have the following data frame.

df

Nodes       Weight
A,B          10
A,C,F        8
B,F,D        6
B,E          4

I would like to split based on comma and keeping their weight as well. For example, Nodes (A,C,F) A has connection with C and C has connection F. So, I would like to see A >>C, and C>>F. No need to see A>>F. AND their weight should be 8 as well as shown below.

The final data frame that am looking for looks like below.

Node_1    Node_2     Weight
A            B        10
A            C        8
C            F        8
B            F        6
F            D        6
B            E        4

The goal of creating this data frame is creating a network graph out of it.

There are similar solutions but I couldn't get the result what I want.

I tried with the following:

df = (df['Nodes'].str.split(',') .groupby(df['Weight'])

Can anyone help on this?

2 Answers

Here is one way to do this:

# From https://docs.python.org/3/library/itertools.html#itertools-recipes
from itertools import tee
def pairwise(iterable):
    "s -> (s0,s1), (s1,s2), (s2, s3), ..."
    a, b = tee(iterable)
    next(b, None)
    return zip(a, b)

df['Node_pairs'] = df['Nodes'].str.split(',').apply(lambda x: list(pairwise(x)))
df = df.explode('Node_pairs')
df['Node1'] = df['Node_pairs'].str[0]
df['Node2'] = df['Node_pairs'].str[1]
df

Output:

   Nodes  Weight Node_pairs Node1 Node2
0    A,B      10     (A, B)     A     B
1  A,C,F       8     (A, C)     A     C
1  A,C,F       8     (C, F)     C     F
2  B,F,D       6     (B, F)     B     F
2  B,F,D       6     (F, D)     F     D
3    B,E       4     (B, E)     B     E

Details:

  • Use pairwise recipe from itertools documentation to create 'Node_pairs'
  • Explode the dataframe on list of 'Node_pairs'
  • Assign 'Node1' and 'Node2' using .str get shortcut.

The logic is same as the solution provided by the Scott.

def grouper(input_list, n = 2):
    for i in range(len(input_list) - (n - 1)):
        yield input_list[i:i+n]
        
(df.set_index('Weight')['Nodes']
   .str.split(',')
   .map(grouper)
   .map(list)
   .explode()
   .apply(pd.Series).add_prefix('Node_')
   .reset_index())

    Weight  Node_0  Node_1
0   10      A       B
1   8       A       C
2   8       C       F
3   6       B       F
4   6       F       D
5   4       B       E
Related