AtCoder Regular Contest 178 C - Sum of Abs 2

競プロ AtCoder ARC DP

問題

15 分くらい間に合わなくて悔しい。

解法

以下 0-indexed であるとする。

\sum_{j=0}^{L-2} \sum_{k=j+1}^{L-1} |B_j - B_i| = f(B) とする。

まず数列 B はソートしても良い。 また、相対的な差にしか興味がないため、 B_0 = 0 であるとして良い。

数列の差分を図示して、各差分の寄与を B_{L - 2} から逆順に考えると以下のようになる。（以下は L = 5 の場合である。）

数列のある隣接要素の差分を 1 増やすと、その差分の寄与だけ f(B) が増加する。（例えば L = 5 なら寄与は 4 (=4,1+1+1+1) または 6 (=3+3,2+2+2) のいずれかとなる。）

この寄与は上記図の計算方法からわかるように、差分 B_{k+1} - B_{k} は L-k-1 が k+1 回だけ足されているため (k+1)(L-k-1) (k=0,1,...,L-2) と表現することができる。

B はソートされているため、 \max B = B_{L-1} であり、 B_{L-1} は差分の総和である。

最終的には B_{L-1} を最小化するため、 f(B) = A_{i} となるように隣接要素の差分を 1 増やす操作の回数を最小化すると言い換えることができ、これはナップザック問題を解けば良い。

\max A = M とすると、試すべき k は O(L) 個あるため 全体の計算量は一見 O(ML) となりそうだが、実は試すべき k は O(\sqrt{M}) 個で抑えられるため O(M \sqrt{M}) である。

理由

(k+1)(L-k-1) \leq M となる k についてのみ考えれば良い。

L^2 \leq 4M のとき、 L \leq 2\sqrt{M} より k は O(\sqrt{M}) 個。

L^2 \gt 4M のとき、 (k+1)(L-k-1) \leq M を解くと

k \leq \frac{L-2 - \sqrt{L^2 - 4 M}}{2}, \frac{L-2 + \sqrt{L^2 - 4 M}}{2} \leq k で、対称性より k \leq \frac{L-2 - \sqrt{L^2 - 4 M}}{2} となる k だけ見れば良い。

ここで

\frac{L-2 - \sqrt{L^2 - 4 M}}{2}

= \frac{\sqrt{L^2} - \sqrt{L^2 - 4 M}}{2} - 1 \left(\because L \gt 0 \right)

\leq \frac{\sqrt{L^2 - (L^2 - 4 M)}}{2} - 1 \left(\because \sqrt{a} - \sqrt{b} \leq \sqrt{a - b} (a \geq b \geq 0)\right)

= \frac{2\sqrt{M}}{2} - 1

= \sqrt{M} - 1

なので k は O(\sqrt{M}) 個。

#include "my_template.hpp"
using namespace std;

void solve() {
    INT(N, L);
    VEC(int, A, N);
    constexpr int M = 200000;
    vector<int> dp(M + 1, INF<int>);
    dp[0] = 0;
    REP(k, L - 1) {
        i64 w = (k + 1) * (L - 1 - k);
        REP(j, M + 1 - w) chmin(dp[j + w], dp[j] + 1);
    }
    REP(i, N) {
        int ans = dp[A[i]];
        if (ans == INF<int>) ans = -1;
        print(ans);
    }
    return;
}

int main() {
    solve();
    return 0;
}

感想

B がソート済みのとき、 \sum_{j=0}^{L-2} \sum_{k=j+1}^{L-1} |B_j - B_i| = \sum_{k=0}^{L - 2} (k+1)(L-k-1) (B_{k+1} - B_{k}) というのがかなり本質な気がする。

ChatGPT でこの式変形を導出できた人のツイートが TL に流れてきてびっくりした。（自分のChatGPTはまるで役に立たなかった。）