トコロテンの日記

自分の活動内容や普段考えていることをアウトプットするためのブログです。ITに関わる話がメインであり、開発、競技プログラミング、気に入った技術などの話が多くなります。

AtCoder Beginner Contest 102 C - Linear Approximation

競技プログラミング AtCoder ABC

問題

atcoder.jp

考察と解法

今回は結構考察パートが長めです。

問題としては以下の関数f(A, b)の値を最小化するbを求め、それを代入したf(A, b)の値を求めるといったものです。 A = (A_1, A_2, ..., A_N)は入力として与えられるのでbの値だけを考えれば良いです。

f(A, b) = abs(A_1 - (b + 1)) + abs(A_2 - (b + 2)) + ... + abs(A_N - (b + N))

まず、(-1, -2, ..., -N)が定数になるため、これを事実上の定数である(A_1, A_2, ..., A_N)と合わせて、式変形を行います。

f(A, b) = abs(A_1 - (b + 1)) + abs(A_2 - (b + 2)) + ... + abs(A_N - (b + N)) f(A, b) = abs(A_1 - 1 - b) + abs(A_2 - 2 - b) + ... + abs(A_N - N - b)

ここでA' = (A_1', A_2', ..., A_N') = (A_1 - 1, A_2 - 2, ..., A_N - N)とすると

f(A, b) = abs(A_1' - b) + abs(A_2' - b) + ... + abs(A_N' - b) = \sum_{i=1}^N abs(A_i' - b)

となります。これは、全てのA_i'とbとの距離の総和を意味しています。bは自明に(min(A_i') \le b \le max(A_i'))になりますが、A_i'の取る値が1から10^9であるため、いつも通り全探索では時間内に計算が終わりません。おとなしく問題と向き合いましょう。
また、ここで考えるのを簡単にするために、A'の各要素を昇順(小さい順)に並べ替えたものをBとします。 これはもちろん

f(A, b) = \sum_{i=1}^N abs(B_i - b)

が成り立つため、全てのB_iとbの距離の総和について考えることにします。

ここで、あるb_iをjだけ増加させた場合、f(B, b_i + j)は以下のようになります。

B_i \le b_i \le B_{i+1}(1 \le i \le N - 1)

({B_i - b_i} \le j \le {B_{i+1} - b_i})

f(B, b_i + j) = f(B, b_i) + ij - (N - i)j
f(B, b_i + j) = f(B, b_i) + (i - (N - i))j

したがって、bをjだけ増加させた場合の関数の増加量は

f(B, b_i + j) - f(B, b_i) = (i - (N - i))j (増加式)

となります。f(B, b_i)を最小化するためには、上式を最小化すればよいです。
右辺に着目すると

(i - (N - i))j = (2i - N)j

上の式が最小となるのは以下の2パターンです。

  1. (2i - N) \ge 0かつ(j = B_i - b_i)
  2. (2i - N) \le 0かつ(j = B_{i+1} - b_i)

言い換えると、jの値を

  1. (i \ge N - i)のとき(j = B_i - b_i)
  2. (i \le N - i)のとき(j = B_{i+1} - b_i)

とすることで関数の増加量を最小化できます。また、最小化した場合、(増加式)が0以下の値を取るため、f(B, b_i) \ge f(B, b_i + j)が成り立ち、f(B, b_i)の最小値はf(B, b_i + j)となります。
そして、f(B, b_i + j)の最小値、つまりf(B, b_i)の最小値は

  1. (i \ge N - i)のときf(B, b + j) = f(B, B_i)
  2. (i \le N - i)のときf(B, b + j) = f(B, B_{i+1})

となります。
また、f(B, b)の最小値は自明に全てのi(1 \le i \le N - 1)に対するf(B, b_i)の最小値であることから、

min(f(B, b)) = min(f(B, b_1), f(B, b_2), ..., f(B, b_{N - 1}))

であり、b_iはiによってB_i, B_{i+1}の適切な方を選べばよいことがわかったため、

min(f(B, b)) = min(f(B, B_1), f(B, B_2), ..., f(B, B_N))

となります。あとは、各iについてf(B, B_i)を求める方法がわかればこの問題を解く事ができます。

ここで、f(B, B_i)の式を確認すると

f(B, B_i) = abs(B_1 - B_i) + abs(B_2 - B_i) + ... + abs(B_i - B_i) + ... + abs(B_N - B_i)

abs(B_i - B_i) = 0より

f(B, B_i) = \sum_{j=1}^{i-1} abs(B_j - B_i) + \sum_{j=i+1}^N abs(B_j - B_i)

Bが昇順に並んでいるという制約から

  1. (j \le i)ならabs(B_j - B_i) = B_i - B_j
  2. (j > i)ならabs(B_j - B_i) = -B_i + B_j

よって、

f(B, B_i) = \sum_{j=1}^{i-1} (B_i - B_j) + \sum_{j=i+1}^N (-B_i + B_j)

f(B, B_i) = B_i(i - 1) - \sum_{j=1}^{i-1} B_j + \sum_{j=i+1}^N B_j - B_i(N - i)

f(B, B_i) = B_i(2i - N - 1) - \sum_{j=1}^{i-1} B_j + \sum_{j=i+1}^N B_j (一般式)

となります。\sum_{j=1}^{i-1} B_jと\sum_{j=i+1}^N B_jは予め累積和を用いて計算しておく事で各iに対してO(1)で求める事ができます。

以上の(一般式)を全てのiに対して計算した結果の最小値がこの問題の解となります。

計算量

  1. 各要素をソートするのにO(NlogN)
  2. 累積和を計算するのにO(N)
  3. 全てのiに対して(一般式)を計算するのにO(N)

1, 2, 3の操作は互いに独立しているため、全体での時間計算量はO(NlogN)になります。

実装の注意点

この問題では、(一般式)が取りうる最大値を取った時に大体2 \times 10^{14}となるため、32ビットで表現できる数値の範囲を超えてしまいます。 したがって、殆どの環境で32ビットとなるint型ではオーバーフローが発生します。C++ではlong long型が64ビットの幅を持っているため、これを利用しましょう。

ソースコード

#include <iostream>
#include <algorithm>
 
#define int long long
 
using namespace std;
 
signed main()
{
    int n;
    int a[200000];
    int sum[200001] = { 0 };
 
    cin >> n;
 
    for(int i = 0; i < n; ++i)
    {
        cin >> a[i];
        // 予めAをA'に変換しておく
        a[i] -= i + 1;
    }
 
    // A'からBへの変換
    sort(a, a + n);
 
    // 累積和を計算[)
    for(int i = 1; i <= n; ++i)
    {
        sum[i] = sum[i - 1] + a[i - 1];
    }

    // f(B, b)の最小値
    int d = 1LL << 60LL;

    for(int i = 0; i < n; ++i)
    {
        int b = a[i];
        int leftD = (i * b) - (sum[i]);
        int rightD = (sum[n] - sum[i + 1]) - (n - 1 - i) * b;
        d = min(d, leftD + rightD);
    }
 
    cout << d << endl;
}

感想

300点の問題にしては少し難しいと感じました。bの値の候補をいかに絞り込む事ができるかがこの問題の鍵でした。