python条件函数: np.where 和 np.select

原创

极致python

发布于 2025-03-19 15:13:15

85800

代码可运行

运行总次数：0

代码可运行

多条件选择：根据多个布尔数组选择数据, 根据不同条件将数据选择不同的类别

一、np.where

1、返回满足条件的元素的索引

import numpy as np
arr=np.array([0, 1, 2, 3, 4, 5, 6])
#找到数组中所有大于3的元素的索引
indices=np.where(arr>3)
print(indices)
#输出结果: (array([4, 5, 6]),)

2、基于一个条件在两个数组之间选择元素

import numpy as np
x=np.array([1, 2, 3, 4, 5])
y=np.array([10, 20, 30, 40, 50])
#根据条件选择元素，如果条件为真则选择x的元素，
#否则选择y的元素
result=np.where(x<3, x, y)
print(result)
#输出结果：[1 2 30 40 50]

score_array=np.array([56, 61, 95])
score_result=np.where(
        score_array>=60, 
           "及格",  "不及格"
        )
print(score_result)
#输出结果: ['不及格' '及格' '及格']

3、根据DataFrane中其他列的条件添加新列

import pandas as pd
import numpy as np
df=pd.DataFrame({
      'A':[4, 5, 9, 6, 5],
      'B':[4, 7, 3, 8, 5],
      'C':[1, 9, 5, 2, 5],
      'D':[2, 1, 4, 8, 9]
       })
#新建E列，如果A列中大于4且D列小于8，
#则E列为True,否则为False
df['E']=np.where(
      (df['A']>3)&(df['D']<8), 
      True, False
      )
print(df)

输出结果：

二、np.select

np.select(condlist, choicelist, default=0)

根据condlist中每个对应的条件，在choicelist选择对应的返回结果

1、数组

import numpy as np
x=np.arange(6)
condlist=[x<3, x>3]
choicelist=[x, x**2]
result=np.select(condlist, choicelist, 42)
print(result)
#输出结果：[0 1 2 42 16 25]

score_array=np.array([56, 61, 95])
score_result=np.select(
        score_array>=60, 
        ["及格"]*len(score_array), 
        ["不及格"]*len(score_array)
        )
print(score_result)
#输出结果：['及格' '及格' '及格']

2、根据DataFrame中其他列的条件添加新列

import pandas as pd
import numpy as np
df=pd.DataFrame({
      'A':[4, 5, 9, 6, 5],
      'B':[4, 7, 3, 8, 5],
      'C':[1, 9, 5, 2, 5],
      'D':[2, 1, 4, 8, 9]
       })

conditions=[
    (df['D'] <=2),
    ((df['D'] >2) & (df['D'] <=5)),
    ((df['D'] >5) & (df['D'] <=8)),
    (df['D']<=10)
    ]

values=['tier1','tier2','tier3','tier4']
df['tier']=np.select(conditions, values, default="other")
print(df)

输出结果：