找到一个密集向量的最大值将其设置为1,其余所有设置为0 [pyspark]
问题描述:
我有一个包含密集向量的DataFrame列。我想一个这些矢量的最大值设置为1且取代所有其余为0。 例如:找到一个密集向量的最大值将其设置为1,其余所有设置为0 [pyspark]
+------+-------------+
| Col1| vector|
+------+-------------+
|Modali|[1.0,2.0,3.0]|
|assert|[4.0,9.0,3.0]|
+------+-------------+
我想将其变为:
+------+-------------+
| Col1| vector|
+------+-------------+
|Modali|[0.0,0.0,1.0]|
|assert|[0.0,1.0,0.0]|
+------+-------------+
注意我正在使用python开发Spark 2.1.0。 非常感谢你提前。
答
您可以创建一个udf
,它接受一个向量并将其二进制化;可以通过简单地检查向量中的值是否等于使用列表理解的最大值来构造二进制器:
from pyspark.ml.linalg import Vectors, VectorUDT
import pyspark.sql.functions as F
def max_binarizer(vector):
max_val = max(vector) # maximum value in the vector
return Vectors.dense([1 if x == max_val else 0 for x in vector]) # binarize it
# create a udf for the binarizer
max_bin_udf = F.udf(max_binarizer, VectorUDT())
df.withColumn("vector", max_bin_udf(df["vector"])).show()
+------+-------------+
| Col1| vector|
+------+-------------+
|Modali|[0.0,0.0,1.0]|
|assert|[0.0,1.0,0.0]|
+------+-------------+
这工作,非常感谢 – abdelkarim