---
tags: AtCoder
---
# KEYENCE2019-E Connecting Cities を $\Theta(N)$ で

解けます とても良い問題

## 問題文

<https://atcoder.jp/contests/keyence2019/tasks/keyence2019_e>

密なグラフの最小全域木です.

## 考察

全ての辺のコストを求めることはできないので, コスト関数の構造をうまく使うことが必要です.

ここで、辺の片側を $i$ に固定してみましょう. そうするとコスト関数は
$$
\text{cost}(j)=D\times|i-j|+A_j+A_i
$$
になります。 絶対値は展開すると見通しが良くなることがあります. 今回は当てはまるので, 展開しておきましょう.

$$
\text{cost}(j)=
\begin{cases}
-D \times j + A_j + (D \times i + A_i)&(j<i)\\
D \times j + A_j + (-D \times i + A_i)&(j>i)\\
\end{cases}
$$

$i$ の関わるところは定数なのでくくっておきました.
さて, この状態で, $i$ を 1 増やしてみます.

$$
\text{cost}(j)=
\begin{cases}
-D \times j + A_j + (D \times i + D + A_{i+1})&(j<i+1)\\
D \times j + A_j + (-D \times i - D + A_{i+1})&(j>i+1)\\
\end{cases}
$$

よく見ると, 定数部分しか変わっていないですね. つまり, 大小関係は変わっていないということになります.

さて, コスト関数の構造がわかったところで, 最小全域木を作っていきましょう.
少し変わった最小全域木の問題を解くときは, 最小全域木のアルゴリズムがうまく使えないか考えると解けがちです.
プリム法やブルーフカ法に思いを馳せると, 以下の性質があることに気がつきます.

- 全ての頂点を 2 グループに分ける. このとき, 2 グループにまたがる辺の中でコスト最小のものは必ず採用される.(最小のものが複数あるとき, そのうち 1 つに注目すると, その辺を含む最小全域木は存在する.)

[証明]
2 グループにまたがる辺の中でコスト最小のものが 1 つしかない場合のみ証明する.
最小全域木がコスト最小の辺を使わないと仮定する. この最小全域木にコスト最小の辺を加えたなもりグラフを考えると, このグラフはコスト最小の辺を含むサイクルが存在する. このサイクルの中のコスト最大の辺を削除すると, 最小全域木よりコストの小さい木になって, 矛盾である.

---

この性質を使いましょう. コスト関数の性質から, $A[1,i]$ と $A[i+1,N]$ に分けるのが良さそうです.
この 2 つに分けたとき, 2 グループにまたがる辺の中でコスト最小のものはどれになるでしょうか？
大小関係が変わらない性質から, $A[1,i]$ 側は $\displaystyle \underset{j\,\in\,[1,i]}{\arg\min}\{-D \times j + A_j\}$ , $A[i+1,N]$ 側は $\displaystyle \underset{j\,\in\,[i+1,N]}{\arg\min}\{D \times j + A_j\}$ を選べばコスト最小になり, この 2 頂点を結ぶ辺は必ず使われます.

全ての $i$ について $\displaystyle \underset{j\,\in\,[1,i]}{\arg\min}\{-D \times j + A_j\}$ と $\displaystyle \underset{j\,\in\,[i+1,N]}{\arg\min}\{D \times j + A_j\}$ を $\Theta(N)$ で求めて, 全ての辺を採用します.
これらの辺を採用したとき, グラフはどのようになっているでしょうか？

まず, 使われない頂点があることが分かります. そして, $\displaystyle \underset{j\,\in\,[i+1,N]}{\arg\min}\{D \times j + A_j\}=i+1$ のときに $i$ を 1 増やすことを考えると, 使われない頂点以外が木をなすことも分かります.

最後に, 使われない頂点をつなげていけば, 最小全域木の完成です.
使われない頂点 $i$ について, $i$ につながる最もコストの低い辺を考えましょう.
コスト関数の構造から, もう一端は $\displaystyle \underset{j\,\in\,[1,i-1]}{\arg\min}\{-D \times j + A_j\}$ か $\displaystyle \underset{j\,\in\,[i+1,N]}{\arg\min}\{D \times j + A_j\}$ のいずれかになります. 実は, この 2 つはすでに計算済みで, しかも "使われる頂点" のほうなので, 2 つの辺のうちコストの低い方を採用すれば良いです.

これより, $\Theta(N)$ でこの問題を解くことができました.

## 実装

https://atcoder.jp/contests/keyence2019/submissions/15891705

#include <bits/stdc++.h>
using namespace std;
using ll = int64_t;
using pll = pair<ll, ll>;


int main(){
    ll n, d;
    cin >> n >> d;
    vector<ll> a(n);
    for(auto& i : a) cin >> i;
    auto cost = [&](ll x, ll y){
        return a[x] + a[y] + abs(x - y) * d;
    };
    vector<ll> l(n), r(n);
    l[0] = 0;
    for(ll i = 1; i < n; i++) l[i] = min(l[i - 1], i, [&](ll x, ll y){ return cost(i, x) < cost(i, y); });
    r[n - 1] = n - 1;
    for(ll i = n - 1; i--; ) r[i] = min(r[i + 1], i, [&](ll x, ll y){ return cost(i, x) < cost(i, y); });
    ll ans = 0;
    pll last;
    for(ll i = 0; i < n - 1; i++) if(last != pll{l[i], r[i + 1]}){
        last = {l[i], r[i + 1]};
        ans += cost(l[i], r[i + 1]);
    }
    for(ll i = 1; i < n - 1; i++) if(l[i] != i && r[i] != i) ans += min(cost(i, l[i - 1]), cost(i, r[i + 1]));
    cout << ans << endl;
}