AT274の技術だったりなかったり

あっとのTECH LOG

競プロのこととか技術のこととか。たまに日常を書きます。

ABC151 E - Max-Min Sums

問題原文

atcoder.jp

問題要旨

有限個の整数からなる集合  X に対し、  f(X) = maxX - minX とする。
 N 個の整数  A_1, A_2 ... , A_N が与えられる。このうちから  K 個選び、それを  X とする。 全ての K 個 の選び方について、  f(X) を計算し、その総和を  10^{9} + 7 で割ったあまりを求めよ。

解法

こういうのは、「各数が最大値/最小値としてカウントされるような選び方は何通りあるか?」を考えればいい。
ある  A_i が最小値になるような選び方を考えた時に、 ある  A_j が最大値になるような選び方を考えていることにもなるが、 最小値と最大値のパターンを別個に計算すれば問題ない。重複があってもOK。

実装

途中で大きい順に並び替えれば、最大値については最小値のコードのコピーでいける。

N, K = map(int, input().split())
A = sorted(list(map(int, input().split())))
MOD = 10 ** 9 + 7
 
ans = 0
# 各Aiについて、それが最小値になるような選び方が何通りあるかを求めながら計算
for i in range(N - K + 1):
    ans -= A[i] * nCr(N - i - 1, K - 1) % MOD
    ans %= MOD
 
# 各Aiについて、それが最小値になるような選び方が何通りあるかを求めながら計算
A = A[::-1]  # 大きい順に並び替え!
for i in range(N - K + 1):
    ans += A[i] * nCr(N - i - 1, K - 1) % MOD
    ans %= MOD
 
print(ans % MOD)

感想

9分で解いたらしい。えらい。
- 求めたいものを分けて考える - 各要素について、何回数えられるかを考える

考え方はド典型なので、常に引き出せるようにしたいですね。