找到一個密集向量的最大值將其設置爲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]| 
+------+-------------+