ABC149 E - Handshake
問題
$ N個からなる数列$ A_iが与えられます。
2つ組$ (x_i, y_i) (0 \leq x_i \leq N-1, 0 \leq y_i \leq N-1)を重複しないように$ M個選びます。
$ \sum_{i=0}^{M-1} A_{x_i} + A_{y_i}の最大値はいくらでしょう。
制約
$ 1 \leq N \leq 10^5
$ 1 \leq M \leq N^2
$ 1 \leq A_i \leq 10^5
入力はすべて整数
考察
愚直に和を計算すると$ O(N^2)かかってしまいます。
まず、簡単な場合を観察してみます。
$ Aはソートしておきます。
$ M = 1の場合は明らかに$ (N-1, N-1)です。
$ M = 3の場合までは簡単で$ (N-1,N-1), (N-1, N-2), (N-2,N-1)を選べばよいです。
$ M = 4では$ (N-1, i), (N-2, N-2)のどちらが大きいかわからないので難しくなります。ただ、この場合でも$ A_{N-2} + A_{N-2} - A_{N-1}で二分探索すれば$ (N-2,N-2)より大きい$ (N-1,i)がいくつあるか$ O(\log N)でわかります。 これ以降はかなり複雑で一般的に考えるのが難しそうです。
もし、これ以上の大きいものを探したいというハードル$ tがわかっているなら、$ M = 4の場合のように$ k - A_iで二分探索することで$ iを使って和が$ t以上になる個数がわかります。すべての$ iについて試すことでトータルで$ O(N\log N)で和が$ t以上になるものの個数がわかります。この個数が$ Mをまたぐ境界の$ tを二分探索で探しましょう。 二分探索によって、使う個数が$ M以下となる最小の$ t($ t^*としておきます)を探します。 それぞれの$ iについて
二分探索で、和が$ t - A_i以上となる最小のインデックス$ jを探すと$ N - j個が$ t以上になるものだということがわかります。 $ t^*が与えられたときの、和が$ t^*以上になるものの和は、$ t^*を探すときと似たように
それぞれの$ iについて
二分探索で、和が$ t^* - A_i以上となる最小のインデックス$ jを探す 和は$ (A_i + A_j) + (A_i + A_{j+1}) + \cdots (A_i + A_{N-1})
$ = (N-j)A_i + \sum_{l=j}^{N-1} A_l
なので、累積和を使うと$ O(1)で計算できます。 $ t^*以上のものを使ったときにちょうど$ M個になればそれで終わりですが、$ Mより小さかった場合にはどうすればよいでしょう?
$ t^* -1以上のものを選ぶと$ M個を越えてしまうのですから、和が$ t^* - 1となるものによって必ず$ M個に足りない分を使い切ることが出来ます。
時間計算量は$ O(N \log N \log A_{\max})です。
細かい部分を詰めるのに結構手間取りました。
code:cpp
using namespace std;
typedef long long ll;
#define rep(i, N) for (int i = 0; i < (int)N; ++i) #define all(a) (a).begin(), (a).end() ll N;
ll M;
vector<ll> A;
int main() {
cin >> N >> M;
A.resize(N);
sort(all(A));
vector<ll> S(N + 1, 0);
function<ll(ll)> num = &(ll t) { ll res = 0;
rep(i, N) res += distance(lower_bound(all(A), t - Ai), A.end()); return res;
};
ll l = 0, r = 1e6;
while (r - l > 1) {
ll m = (l + r) / 2;
if (num(m) <= M) {
r = m;
} else {
l = m;
}
}
ll ans = 0;
ll rest = M;
rep(i, N) {
int l = distance(lower_bound(all(A), r - Ai), A.end()); rest -= l;
}
ans += rest * l;
cout << ans << endl;
return 0;
}
$ A_i = xとなるものの度数を$ F \lbrack x \rbrackとしたとき、$ A_i + A_jの度数$ F'は、
$ F'\lbrack x \rbrack = \sum_{y=0}^{x} F[y]F[x-y]
したがって高速フーリエ変換で$ O(A_{\max} \log A_{\max})で解くことが出来ます。 たたみ込みを計算するライブラリさえ持っていればすごくシンプルに解けますね!
code:cpp
using namespace std;
typedef long long ll;
#define rep(i, N) for (int i = 0; i < (int)N; ++i) #define all(a) (a).begin(), (a).end() int N;
ll M;
vector<int> A;
int main() {
cin >> N >> M;
A.resize(N);
const int A_MAX = *max_element(all(A));
vector<ll> freq(A_MAX + 1, 0);
Convolution<ll> conv(freq, freq);
ll ans = 0;
for (int i = 2 * A_MAX; i >= 0; --i) {
ans += c * i;
M -= c;
if (M == 0) break;
}
cout << ans << endl;
return 0;
}
Convolutionの実装を見たい場合は以下の提出コードを見てください。