ログインしてください。
D - 距離 (Distance) 解説
by
seekworser
O(N)解法
\(O(N)\) でこの問題を解く方法を解説します。(発展的な解法です。)
ある色 \(X\) を決めて、色が \(X\) であるような添字を昇順に並べた列を \(A = (A_1, A_2, \dots, A_M)\) とします。 このとき、ある \(A_k\) に対する答えは、
\[ \begin{aligned} \sum_{i=1}^{m} | A_i - A_k |=& \sum_{i=1}^{k-1} (A_k - A_i) + \sum_{i=k}^{M} (A_i - A_k)\\ =& (k - 1) A_k - \sum_{i=1}^{k-1} A_i + \sum_{i=k}^{M} A_i - (M - k + 1) A_k \end{aligned} \]
です。(添字は昇順に並んでいることを利用して変形しています。) 2つの総和の箇所については、\(k\) を昇順に走査していけば \(O(1)\) で差分更新ができるので、配列全体に対して \(O(M)\) で答えを得ることができます。
各色ごとに添字の配列を作成することで、この問題を全体で \(O(N)\) で解くことができます。
実装例:(以下の実装例では、Aの添字が解説と異なり 0-indexed であることに注意してください。)
#include <bits/stdc++.h>
using namespace std;
int main() {
int n; cin >> n;
vector<int> a(n);
for (size_t i=0; i<n; i++) cin >> a[i];
vector pos(101, vector<int>());
for (int i=0; i<n; i++) pos[a[i]].emplace_back(i);
vector<int> ans(n);
for (size_t i=0; i<pos.size(); i++) {
int lower = 0, upper = 0;
for (auto x : pos[i]) upper += x;
for (size_t j=0; j<pos[i].size(); j++) {
int x = pos[i][j];
ans[x] = (x * j - lower) + (upper - x * (pos[i].size() - j));
lower += x;
upper -= x;
}
}
for (auto x : ans) cout << x << '\n';
}
投稿日時:
最終更新:
