AtCoder Beginner Contest 143 に参加しなかったので、復習して ABCD まで解いた。解説にあった B を O(N) で解く方法を考えた。
以下のコードは C++14 (GCC 5.4.1) で AC を確認している。
問題 B
HP の回復量は愚直に書き下すと下記のように \(O(N^2)\) になる。
$$ \sum_{i = 1}^{N - 1} \sum_{j = i + 1}^{N} d_i \times d_j $$
最内部の \(d_i\) は \(j\) に依存しないので \(j\) に関する和の外に出し
$$ \sum_{i = 1}^{N - 1} d_i \times \sum_{j = i + 1}^{N} d_j $$
と書き換えても良い。つまり \(j\) に関する部分和の配列を \(O(N)\) で計算し、その部分和と \(d_i\) の積の総和を求める計算を \(O(N)\) で計算することで、全体として \(O(N)\) で計算できる。
#include <bits/stdc++.h>
using namespace std;
int main()
{
int N;
cin >> N;
vector<int> d;
d.resize(N);
for (int i = 0; i < N; ++i) {
cin >> d[i];
}
vector<int> d_sum;
d_sum.resize(N - 1);
// d に関する部分和は、最後の要素一つのものから順に計算していく
d_sum[0] = d[N - 1];
for (int i = 1; i < N - 1; ++i) {
d_sum[i] = d_sum[i - 1] + d[N - i - 1];
}
// d_sum = {
// d[N - 1]
// d[N - 1] + d[N - 2]
// ...
// d[N - 1] + d[N - 2] + ... + d[1]
// }
uint64_t ans = 0;
for (int i = 0; i < N - 1; ++i) {
ans += d[i] * d_sum[N - i - 2];
}
// ans =
// d[0] * (d[N - 1] + d[N - 2] + ... + d[1])
// + ...
// + d[N - 3] * (d[N - 1] + d[N - 2])
// + d[N - 2] * d[N - 1]
// }
cout << ans << endl;
return 0;
}