E - Chinese Restaurant (Three-Star Version) Editorial by sounansya


\(a[i]\)\(i\) 回の操作で不満度が \(0\) になるような人の数 と定義します。(これは \(p\) から簡単に求めることができます。) すると、 \(i\) 回操作した後の \(N\) 人の不満度の総和は \(\displaystyle \sum_{k=0}^{N-1}a[(k-i)\% N]\min(k,N-k)\) となります。 \(b[k]=\min(k,N-k)\) とすることでこの総和は \(\displaystyle \sum_{k=0}^{N-1}a[(k-i)\% N]b[k]\) と書くことができます。

ここで \(b[k]=b[N-k]\) であることに着目すると、

\[\begin{aligned} \sum_{k=0}^{N-1}a[(k-i)\% N]b[k]=\sum_{k_1+k_2=2N-i}a[k_1]b[k_2]+\sum_{k_1+k_2=N-i}a[k_1]b[k_2] \end{aligned}\]

と表せることが分かります。 \(c\)\(a,b\) の畳み込みとするとこれは \(c[2N-i]+c[N-i]\) と書くことができます。

以上により求める値は \(\displaystyle \min_{0\le i<N}(c[2N-i]+c[N-i]) = \min_{1\le i\le N}(c[i]+c[N+i])\) となります。計算量は畳み込みがネックとなり \(O(N\log N)\) です。公式解説と比べ計算量は劣りますが、こちらの方がバグを埋め込みにくいです。

畳み込みの実装によっては \(c[2N]\) が存在しないことに注意してください。

実装例(C++)

#include <bits/stdc++.h>
#include <atcoder/convolution>
using namespace std;
using ll = long long;
#define rep(i, n) for (ll i = 0; i < n; i++)
int main()
{
    int n;
    cin >> n;
    vector<int> p(n);
    rep(i, n) cin >> p[i];
    vector<ll> a(n), b(n);
    rep(i, n)
    {
        int idx = p[i] - i;
        if (idx < 0)
            idx += n;
        a[idx]++;
        b[i] = min(i, n - i);
    }
    vector<ll> c = atcoder::convolution_ll(a, b);
    ll ans = 1e18;
    rep(i, n)
    {
        ll tmp = c[i];
        if (i != n - 1)
            tmp += c[n + i];
        ans = min(ans, tmp);
    }
    cout << ans << '\n';
    return 0;
}

posted:
last update: