Can a list comprehension be divided in two lists?-CodePudding

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

CodePudding user response：

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])

CodePudding user response：

Right now your code is overwriting what's in Column A.

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

The above answer is much better than mine. Learned something new today.