けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AOJ 2955 Two Colors Sort (HUPC 2019 day2-D)

これ好き!!!

問題概要

 1, 2, \dots, N の順列が与えられる。これらの順列のうち、 R 個を赤く塗り、 N-R 個を青く塗る方法であって、次の条件を満たすようにするものは存在するか、判定せよ。

  • 赤い要素同士を swap することを好きな順序で好きな回数だけ繰り返し
  • 青い要素同士を swap することを好きな順序で好きな回数だけ繰り返すことで
  • 順列全体がソートされた状態にすることができる

制約

  •  1 \le N \le 3 \times 10^{5}

考えたこと

つまり、

  • 赤く塗った要素の index (1-indexed) の集合
  • 赤く塗った要素の集合

が一致するようにできるかどうかを問う問題となっている。具体的に様子を掴んでみる。たとえば  p_{1} = 3 となっていたとき、もし  p_{1} を赤く塗るならば、 p_{3} も赤く塗る必要があることがわかる。一般に

  •  p_{i} を赤く塗るならば、 p_{p_{i}} も赤く塗る必要がある

ということが言える。この調子で行くと、順列を「巡回群に分ける」という考察が生まれてくる。同一の巡回群の各要素は同じ色で塗らなければならないことがわかる。逆に、同一の巡回群の各要素が同じ色で塗られているならば、目的の操作は達成できる!!よって問題は以下のものに帰着された


順列の巡回群のサイズを  a_{1}, a_{2}, \dots, a_{K} とする。これらの部分集合であって総和が  R となるものが存在するかどうかを判定せよ。


DP へ

部分和問題になったので、まずは次の DP をしたくなる。

  • dp[ i ][ v ] := i 個見て、その部分集合の総和が v となることは可能か?

しかしこのままでは  O(NR) の計算量が必要となって間に合わない。そこで、 a_{1}, \dots, a_{K} をサイズごとに分類することにする。このとき

 a_{1}, \dots, a_{K} の中に異なる値の種類は  O(\sqrt{N}) に抑えられる」

ということが言える。 a_{1} + \dots + a_{K} = N から言える。

  •  a_{i} \le \sqrt{N} であるような  a_{i} はもとより  O(\sqrt{N}) 種類
  •  a_{i} \gt \sqrt{N} であるように  a_{i} は、総和が  N というところから  O(\sqrt{N}) 個以下しかない

という感じ。よって、個数制限付き部分和問題を解けば OK。計算量は  O(N\sqrt{N}) となる。

  • dp[ i ][ v ] := i 種類目までを見て総和を v とする方法のうち、最後の要素を使う個数の最小値

とすれば OK。

コード

#include <bits/stdc++.h>
using namespace std;
template<class T> inline bool chmax(T& a, T b) { if (a < b) { a = b; return 1; } return 0; }
template<class T> inline bool chmin(T& a, T b) { if (a > b) { a = b; return 1; } return 0; }

int main() {
    int N, R;
    cin >> N >> R;
    vector<int> P(N);
    for (int i = 0; i < N; ++i) cin >> P[i], --P[i];

    vector<bool> seen(N, false);
    map<int,int> ma;
    for (int i = 0; i < N; ++i) {
        if (seen[i]) continue;
        int cnt = 0;
        int j = i;
        do {
            seen[j] = true;
            ++cnt;
            j = P[j];
        } while (j != i);
        ma[cnt]++;
    }
    
    int K = ma.size();
    const int INF = 1<<29;
    vector<int> dp(R+1, INF), ndp(R+1, INF);
    dp[0] = 0;
    for (auto it : ma) {
        ndp.assign(R+1, INF);
        int val = it.first, num = it.second;
        for (int v = 0; v <= R; ++v) {
            if (dp[v] < INF) chmin(ndp[v], 0);
            if (v >= val && ndp[v-val] < num) chmin(ndp[v], ndp[v-val]+1);
        }
        swap(dp, ndp);
    }
    if (dp[R] < INF) cout << "Yes" << endl;
    else cout << "No" << endl;
}