pythonpandasdataframelist-comprehensionone-liner

Can a list comprehension be divided in two lists?

I think I've caught the idea of one-line for loop, but now I have a problem. I know I can define a dataframe column using this like:

df = pd.DataFrame(columns=["columnA"])

list = [0, 1, 2, 3, 4]

df["columnA"] = [i for i in list]

Now my question is: Is it possible to define 2 columns in a one-line for loop?

I've tried this:

df["columnA"], df["columnB"] = [i, i**2 for i in list]
df["columnA"], df["columnB"] = [[i, i**2] for i in list]

None of this worked. I'm using Python 3.10

Solution

  • You have to zip your output:

    df['A'], df['B'] = zip(*[(i, i**2) for i in lst])
print(df)

# Output
   A   B
0  0   0
1  1   1
2  2   4
3  3   9
4  4  16

    You can also use np.array:

    df[['A', 'B']] = np.array([(i, i**2) for i in lst])