借助Numpy,优化Pandas的条件检索代码
Numpy
其实是最早的处理数据的Python
库,它的核心ndarray
对象,是一个高效的n维数组结构。
通过这个库,可以高效的完成向量和矩阵运算,由于其出色的性能,很多其他的数据分析,科学计算或者机器学习相关的Python
库都或多或少的依赖于它。
Pandas
就是其中之一,Pandas
充分利用了NumPy
的数组运算功能,使得数据处理和分析更加高效。
比如,Pandas
中最重要的两个数据结构Series
和DataFrame
在内部就使用了NumPy
的ndarray
来存储数据。
在使用Pandas
进行数据分析的过程中,按条件检索和过滤数据是最频繁的操作。
本文介绍两种通过结合Numpy
,一方面让Pandas的检索过滤代码更加简洁易懂,另一方面还能保障检索过滤的高性能。
1. 准备数据
第一步,先准备数据,这次使用二手房交易数据,可从 https://databook.top/lianjia/nj 下载。
import pandas as pd
import numpy as np
# 这个路径替换成自己的路径
fp = r'D:\data\南京二手房交易\南京江宁区.csv'
df = pd.read_csv(fp)
df.head()
2. 一般条件判断(np.where)
比如,买房前我们想先分析下已有的成交信息,对于房价能有个大致的印象。
下面,按照总价和单价,先挑选总价200~300万之间,或者单价1万以下的成交信息。
符合条件返回“OK”,否则返回“NG”。
def filter_data(row):
if row["totalPrice"] > 200 and row["totalPrice"] < 300:
return "OK"
if row["unitPrice"] < 10000:
return "OK"
return "NG"
df["评估"] = df.apply(filter_data, axis=1)
df[df["评估"] == "OK"].head()
上面的过滤数据写法是使用Pandas
时用的比较多的方式,也就是将过滤条件封装到一个自定义函数(filter_data
)中,然后通过 apply
函数来完成数据过滤。
下面我们用Numpy
的 np.where
接口来改造上面的代码。np.where
类似Python
编程语言中的if-else
判断,基本语法:
import numpy as np
np.where(condition[, x, y])
其中:
- condition:条件表达式,返回布尔数组。
- x 和 y:可选参数,
condition
为True
,返回x
,反之,返回y
。
如果未提供x
和 y
,则函数仅返回满足条件的元素的索引。
改造后的代码如下:
# 根据单价过滤
cond_unit_price = np.where(
df["unitPrice"] < 10000,
"OK",
"NG",
)
# 先根据总价过滤,不满足条件再用单价过滤
cond_total_price = np.where(
(df["totalPrice"] > 200) & (df["totalPrice"] < 300),
"OK",
cond_unit_price,
)
df["评估"] = cond_total_price
df[df["评估"] == "OK"].head()
运行之后返回的结果是一样的,但是性能提升很多。
如果数据量是几十万量级的话,你会发现改造之后的代码运行效率提高了几百倍。
3. 复杂多条件判断(np.select)
上面的示例中,判断还比较简单,属于if-else
,也就是是与否的判断。
下面设计一种更复杂的判断,将成交信息评估为“优良中差”4个等级,而不仅仅是“OK”和“NG”。
我们假设:
- 优:房屋精装,且位于中楼层,且近地铁
- 良:总价<300,且近地铁
- 中:总价<400
- 差:其他情况
用传统的方式,同样是封装一个类似filter_data
的函数来判断“优良中差”4个等级,然后用 apply
函数来完成数据过滤。
这里就不演示了,直接看结合Numpy
的np.select
接口,高效的完成“优良中差”4个等级的过滤。
np.select
类似Python
编程语言中的match
匹配,基本语法:
numpy.select(condlist, choicelist, default=0)
其中:
- condlist:条件列表,每个条件都是一个布尔数组。
- choicelist:与 condlist 对应的数组列表,当某个条件为真时,返回该位置对应的数组中的元素。
- default:可选参数,当没有条件为真时返回的默认值。
# 设置 “优,良,中” 的判断条件
conditions = [
df["houseInfo"].str.contains("精装")
& df["positionInfo"].str.contains("中楼层")
& df["advantage"].str.contains("近地铁"),
(df["totalPrice"] < 300) & df["advantage"].str.contains("近地铁"),
df["totalPrice"] < 400,
]
choices = ["优", "良", "中"]
# 默认为 “差”
df["评估"] = np.select(conditions, choices, default="差")
df.head()
这样,就实现了一个对成交信息的分类。
4. 总结
np.where
和 np.select
的底层都是向量化的方式来操作数据,执行效率非常高。
所以,我们在使用Pandas分析数据时,应尽量使用np.where
和 np.select
来帮助我们过滤数据,这样不仅能够让代码更加简洁专业,而且能够极大的提高分析性能。