[Numpy] np.where

用法:满足条件condition,输出x,不满足则输出y 

import pandas as pd
import numpy as np

df = pd.DataFrame(data={'语文':[50,90,70,78,60],
                        '数学':[59,80,60,75,69],
                        '英语':[61,95,65,80,59]},
                  index=['Harry','Andy','Rita','Lee','Jack'])
# 添加'总分'字段
df['总分'] = df['语文'] + df['数学'] + df['英语']
# 调用np.where
# 添加'总分评级'字段
df['总分评级'] = np.where(df['总分']>200,'A等','B等')
# np.where也可以用于两个字段之间的比较
df['语文评级'] = np.where(df['语文'] > df['数学'],'√','×')

df

用法:输出满足条件condition(即非0)元素的坐标

import numpy as np
arr = np.array([2,4,6,8,10])
# 返回索引值2,3,4
np.where(arr > 5) # (array([2, 3, 4], dtype=int64),)
# 等价于arr[arr>5]
arr[np.where(arr > 5)] # array([ 6,  8, 10])

arr

import numpy as np
arr1 = np.array([[0, 1],[1, 0]])
# 输出arr1中值为1的坐标值
# (array([0,1],dtype=int64), array([1,0],dtype=int64))
np.where(arr1)

arr1