forked from andreimaximov/algorithms
-
Notifications
You must be signed in to change notification settings - Fork 2
/
main.cpp
118 lines (101 loc) · 2.43 KB
/
main.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
#include <iostream>
#include <vector>
#include <algorithm>
const uint32_t MOD = 1000000007;
/**
* Calculates a^b % p in logarithmic time
*/
uint64_t modpow(uint64_t a, uint64_t b, uint64_t p) {
uint64_t c = 1;
while (b > 0) {
if (b % 2 != 0) {
c = (c * a) % p;
}
a = (a * a) % p;
b /= 2;
}
return c;
}
/**
* Calculates modular inverse of a mod prime p using Fermat's Little Theorem.
*/
uint64_t modinv(uint64_t a, uint64_t p) {
return modpow(a, p - 2, p);
}
/**
* Calculates nCr % p using the following technique.
*
* C(n, r) = n!/((n-r)!r!)
*
* C(n, r) % p = n! % p * (n-r)!^-1 % p * r!^-1 % p
*
* f(n) = n! % p
* g(n) = n!^-1 % p = (f(n) % p)^-1 % p
*
* C(n, r) % p = (f(n) * g(n-r) * g(r)) % p
*/
template<typename T>
class nCr {
private:
std::vector<T> f;
std::vector<T> g;
T p;
void pre() {
this->f[0] = 1;
for (size_t i = 1; i < this->f.size(); i++) {
this->f[i] = (this->f[i - 1] * i % this->p) % this->p;
}
this->g[0] = 0;
for (size_t i = 1; i < this->g.size(); i++) {
this->g[i] = modinv(this->f[i], this->p);
}
}
public:
nCr(size_t n, T p) :
f(n + 1),
g(n + 1),
p(p) {
this->pre();
}
T calculate(size_t n, size_t r) {
if (r > n) {
return 0;
} else if (n == 0 || r == 0 || r == n) {
return 1;
}
return this->f[n] * this->g[n - r] % this->p * this->g[r] % this->p;
}
};
/**
* For each Xi in sorted nums of size n, we need to calculate a = # of ways
* such that Xi is max of K numbers and b = # of ways such that Xi is the
* minimum of K numbers.
*
* a = C(i, K - 1) because we need to choose K - 1 numbers <= Xi
* b = C(n - i - 1, K - 1) because we need to choose K - 1 numbers >= Xi
*
* We calculate and sum all Xi * (a - b) % p for all Xi in nums to get the
* total.
*/
template<typename T>
uint64_t run(const std::vector<T> &nums, size_t K) {
nCr<uint64_t> ncr(nums.size(), MOD);
int64_t sum = 0;
for (int i = K - 1; i < nums.size(); i++) {
sum += nums[i] * ncr.calculate(i, K - 1) % MOD;
}
for (int i = 0; i < nums.size() - K + 1; i++) {
sum -= nums[i] * ncr.calculate(nums.size() - i - 1, K - 1) % MOD;
}
return (sum % MOD + MOD) % MOD;
}
int main() {
size_t N, K;
std::cin >> N >> K;
std::vector<uint32_t> nums(N);
for (int i = 0; i < N; i++) {
std::cin >> nums[i];
}
std::sort(nums.begin(), nums.end());
std::cout << run<uint32_t>(nums, K) << std::endl;
}