Pyspark获取一列的总和，并使用该值除以另一列

Question

我有一个数据帧df

>>> df = spark.createDataFrame([[1,0], [2,1], [3,1], [4,0], [5,1]], ['a', 'b'])
>>> df.show()
+---+---+
|  a|  b|
+---+---+
|  1|  0|
|  2|  1|
|  3|  1|
|  4|  0|
|  5|  1|
+---+---+

以及

^{pr2}$

使用df，我创建了一个新的数据帧a，它是df的聚合。

>>> a = df.groupby('b').count()
>>> a.show()
+---+-----+
|  b|count|
+---+-----+
|  0|    2|
|  1|    3|
+---+-----+

我需要在a中创建一个名为ev的新列。第i行上的ev的值由

这是我期望的输出

+---+-----+------------------+
|  b|count|           ev_norm|
+---+-----+------------------+
|  0|    2|              1.25|
|  1|    3|0.8333333333333334|
+---+-----+------------------+

但是我首先到达这里，为分子(nrows-count)创建一个新列num，它在两行上产生{}。然后我继续计算分母（denom=0.48），这对所有行都是相同的。最后，我使用.withColumn创建了一个新列，它将num列与{}分开。

有没有一种方法可以动态地执行denom的计算，而不必预先计算它并在一个步骤中执行上述操作？

我试过

a = a.withColumn('ev_norm', (nrows - F.col('count'))/F.sum(F.col('count')*(nrows - F.col('count'))))`

以及

a = a.withColumn('ev_norm', (sum([F.col('count')*(nrows-F.col('count'))]))

但它要么给我一个错误，要么给我一个错误的答案。