如何将列添加到数据框的切片并应用更改

Question

I need to perform a data operation at different subsets of a dataframe and add an additional column, also maintain order of the data.

For example:

df = pd.DataFrame({"x": [1,2,3,4], "y": [0,1,1,0]})

a = df[df.x>2]
b = df[df.x<=2]

a["foo"] = ["a", "b"]
b["foo"] = ["e"]

Now you get the usual warning of slicing operations, dataframe a and b are both changed. However, df remains untouched.

I am wondering if there is a way to get the outcome of the df like:

x     y     foo 
1.    0.    "e"
2.    1.    "e"
3.    1.    "a"
4.    0.    "b"

Note that the operations on setting foo are quite different and complicated, the above is just an example. It's not about how to getting to the end output, more about how to subset a dataframe, perform operations and combine with the additional column added and maintaining the order.

Thanks in advance.

Answer 1

df.loc[df.x.gt(2), 'foo'] = ['a', 'b'] # Note, this only works because df.x.gt(2) returns 2 rows.
df.loc[df.x.le(2), 'foo'] = 'e'

Output:

   x  y foo
0  1  0   e
1  2  1   e
2  3  1   a
3  4  0   b