Branched Evolution

Competitive Programming in Python

ABC 151 E - Max-Min Sums

問題

長さ $n$ の配列 $A$ の長さ $k$ の部分列について、最大値と最小値の差をすべての部分列について足したものを求めよ。($k\leq n\leq 10 ^ 5$)

atcoder.jp

解説

$A$ は昇順にソートしてあるとする。
最小値が $A _ i$ になる長さ $k$ の部分列は $i + 1$ 番以降から $k - 1$ 個選ぶ組合せ分あるので、${} _ {n - (i + 1)} C _ {k - 1}$ 個ある。($i \leq n - k$)
同様に最大値が $A _ i$ になる部分列は ${} _ {i} C _ {k - 1}$ 個ある。($i \geq k - 1$)
以上より、答えは $$ \begin{aligned} & \sum _ {i = k - 1} ^ {n - 1} A _ {i} \times {} _ {i} C _ {k - 1} - \sum _ {i = 0} ^ {n - k} A _ {i} \times {} _ {n - (i + 1)} C _ {k - 1} \\ = & \sum _ {i = k - 1} ^ {n - 1} (A _ {i} - A _ {n - (i + 1)}) \times {} _ {i} C _ {k - 1} \end{aligned} $$

実装

二項係数を何回も計算するときは階乗テーブルと階乗の逆元テーブルを先に計算しておくとよい。
$F _ i =i\ !$, $I _ i =(i\ !) ^ {- 1}$ となるような配列 $F$, $I$ を用意して、 $$ _ {i} C _ {k - 1} = F _ i \times I _ {k - 1} \times I _ {i - ( k -1) } $$ で二項係数を計算する。

n, k = map(int, input().split())
*A, = map(int, input().split())
A.sort()
p = 10**9+7
F, I = [1], [1]
for i in range(1, n+1):
    F.append((F[-1]*i) % p)
    I.append((I[-1]*pow(i, p-2, p)) % p)
s = 0
for i in range(k-1, n):
    s += (A[i]-A[n-(i+1)])*F[i]*I[k-1]*I[i-(k-1)]
    s %= p
print(s)

atcoder.jp