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';
}

投稿日時:
最終更新: