大熊猫栏目的不平等

Question

我有一个pandas数据框,我想基于现有列和某些不等式创建一个新列.例如,让我们

df=pd.DataFrame({'a':[1,2,3,4,5,6,7],'b':[3,6,4,2,7,7,1]})

所以df看起来像

    a   b
0   1   3
1   2   6
2   3   4
3   4   2
4   5   7
5   6   7
6   7   1

我想添加一个新列,res如果相应的值in a小于2,则等于0;如果相应的值in a至少为2且小于4则为1,否则为2.所以我想得到

    a   b   res
0   1   3   0
1   2   6   1
2   3   4   1
3   4   2   2
4   5   7   2
5   6   7   2
6   7   1   2

到目前为止,我一直在使用apply如下:

def f(x):
    if x['a']<2:
        return 0
    elif x['a']>=2 and x['a']<4:
        return 1
    else:
        return 2
df['res']=df.apply(f,axis=1)

但我想知道是否有更直接的方式,或者某些特定的熊猫方法可以让我这样做.

Answer 1

你可以使用pd.cut:

df['res'] = pd.cut(df.a,[-np.inf,2,4,np.inf],labels=[0,1,2],right=False)

输出:

   a  b res
0  1  3   0
1  2  6   1
2  3  4   1
3  4  2   2
4  5  7   2
5  6  7   2
6  7  1   2