首页 > 技术文章 > code_demo 用随机森林做缺失值预测

chenjieyouge 2019-11-05 22:48 原文

直接上代码

在做特征工程的时候, 其实可以用算法来处理特征工程的, 比如缺失值填充之类的. 这里一段code_demo是搬运来的, 不过是真滴好用呢.

# RandomForest - 强化, 对 new_user_class_level 进行 缺失值预测


from pyspark.mllib.regression import LabeledPoint

# 1. 过滤出 new_user_class_level 非空的df 作为 训练集

train_data2 = user_profile_df.dropna(subset=["new_user_class_level"]).rdd.map(
    lambda r:LabeledPoint(r.new_user_class_level - 1, [r.cms_segid, r.cms_group_id, r.final_gender_code, r.age_level, r.shopping_level, r.occupation])
)

# 2. 训练模型

from pyspark.mllib.tree import RandomForest

model2 = RandomForest.trainClassifier(train_data2, 4, {}, 5)

# 3. 测试 一波
model2.predict([0.0, 4.0 ,2.0 , 4.0, 1.0, 0.0])

# 预测值实际应该为2
# 4. 对缺失值进行预测
na_df = user_profile_df.na.fill(-1).where("new_user_class_level=-1")

def row(r):
    return r.cms_segid, r.cms_group_id, r.final_gender_code, r.age_level, r.shopping_level, r.occupation

# 5. rdd方式进行行处理
rdd2 = na_df.rdd.map(row)

# 6. 利用模型预测
predicts = model2.predict(rdd2)

# 7. 修正预测值(+1), 并合并 df 

tmp = predicts.map(lambda x:int(x)).collect()

pdf = pl_na_df.toPandas()

import numpy as np

pdf['new_user_class_level'] = np.array(tmp)+1

# 和预测df 和 真实 df 进行 unionAll
new_user_profile_df = user_profile_df.dropna(subset=["new_user_class_level"]).unionAll(
spark.createDataFrame(tmp, schema=schema))

# 注意: unionAll的使用(纵向合并df), 两个df的表结构必须完全一样
new_user_profile_df.show(5)

推荐阅读