python - 找到密集向量的最大值，将其设置为 1，其余设置为 0 [pyspark]

Question

我有一个包含密集向量的 DataFrame 列。我想将这些向量的最大值设置为 1，并将其余所有向量替换为 0。例如:

+------+-------------+
|  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。提前非常感谢您。

Answer 1

您可以创建一个接收向量并将其二值化的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]|
+------+-------------+