Baby Stepsなブログ

競プロとか。間違ったこと書いてあったら@pi0nep1oneにご連絡ください。

ABC 200 D - Happy Birthday! 2 をDPで解く

atcoder.jp

問題

N個の要素からなる集合A の2つの異なる部分集合で和の mod 200 の値が等しいものが存在するかを判定し、存在する場合はその2つの集合を出力せよ。

制約

 \displaystyle N≦200

公式解説

公式解説は、鳩ノ巣原理により  \displaystyle N≧8 の場合は必ず解が存在することから、 \displaystyle N≦8 の範囲において部分集合を全探索すればよい、というもの。

後にDPによる解法を示すが、明らかにこちらの解法の方が実装上楽である。(DP解法では、テーブルからの数列の復元や、Nが大きい場合に扱う数値が64bit整数型すらオーバーフローする可能性がある等、注意すべき点が多い)

 \displaystyle N≧8 の場合は必ず解が存在する理由は、N要素からなる集合から作れる部分集合の個数は空集合を除いて \displaystyle 2^n - 1 通りなので、 \displaystyle N=8 の時で255通り。この時点で200通りを超えるため、 \displaystyle N≧8 の範囲において mod 200 の値が重複する部分集合が必ず存在すると示せる。(ここが鳩ノ巣原理)

実装上は、 \displaystyle min(N,8) の範囲においてbit全探索を行い、各部分集合が mod 200 のどの値に対応するかを保持し、mod 200 のある値が2通り以上の部分集合で作れるならば、解が存在するとしてその2つの集合を出力すればよい。

実装

#include <bits/stdc++.h>
using namespace std;

#define rep(i, n) for (int i = (int)(0); i < (int)(n); ++i)
#define irep(i, m, n) for (int i = (int)(m); i < (int)(n); ++i)
template<class t> using vc=vector<t>;

signed main()
{
  int n;
  cin >> n;
  vc<int> a(n);
  rep(i, n) cin >> a[i], a[i] %= 200;

  n = min(8, n);
  vc<vc<vc<int>>> p(200);
  irep(s, 1, 1 << n) {
    int sum = 0;
    vc<int> tmp;
    rep(i, n) if ((s >> i) & 1) {
      (sum += a[i]) %= 200;
      tmp.push_back(i);
    }
    p[sum].push_back(tmp);
  }

  rep(i, 200) {
    if (p[i].size() >= 2) {
      cout << "Yes" << endl;
      cout << p[i][0].size() << ' ';
      for (auto e : p[i][0])
        cout << e + 1 << ' ';
      cout << endl;
      cout << p[i][1].size() << ' ';
      for (auto e : p[i][1])
        cout << e + 1 << ' ';
      cout << endl;
      return 0;
    }
  }

  cout << "No" << endl;
}

DP解法

以下の2通りについて解の存在判定の仕方を場合分けする。

  1. Aの全体集合を除く部分集合で和のmod 200が0になるようなものが存在する場合
  2. 1.以外の場合

この様に場合分けするのは、上記の2通りの場合で、部分集合を復元する手順が異なるため。( 1) の場合の集合を 2) 場合の手順で復元しようとするとバグる)

それぞれの方法を以下に示す。

1. 部分集合に 0(mod 200) が存在する場合

(簡単のため、以下に例示する集合Aの要素は全てmod 200を取った値とする)

例として以下の様なものがある。

 \displaystyle A=[101, 1 , 99]

全体集合ではない部分集合の和で 0(mod 200) を作れるならば、その補集合から適当に1要素を取ってきて、Bはその1要素だけ、Cはその1要素と0を作る部分集合とすれば必ず解を作ることができる。

つまり、上の例からB、Cは以下の様になる。

 \displaystyle B=[1]
 \displaystyle C=[101, 1 , 99]

全体集合を除外しているのは、上述の通り 0(mod 200) を作る集合以外に少なくとも1要素は存在しないと、2つの異なる部分集合を作れないため。

2. 1) 以外の場合

1) 以外の場合で解が存在するなら、集合B、Cの総和のmod 200の値は0以外であり、かつその部分集合に 0(mod 200) となるようなものは存在しない。

例として以下の様なもの。

 \displaystyle A=[1, 2 , 3]

 \displaystyle B=[3]
 \displaystyle C=[1, 2]

この場合は以下の様にDPテーブルを定義することで、遷移元を辿ってB、Cを復元することができる。

 \displaystyle dp[i][m]:=i番目の要素までを使って作れる、mod 200の値がmである集合Aの部分集合の個数

 \displaystyle 初期化:dp[0][0]=1
 \displaystyle 遷移:dp[i+1][(m+a[i])\%200] += dp[i][m]

上記から今 dp[i][m] が1以上で、dp[i-1][(m-a[i]+200)%200] も1以上なら総和が m(mod 200) となる部分集合の1つに要素 a[i] を含む者が存在する、ということがわかる。

DPテーブルを更新する過程で、あるテーブルの値が2以上になった瞬間、その m(mod 200) の値を作る部分集合が2つ存在し、かつ、その2つの集合は今見ている a[i] を含むものと、含まないもので2つ存在しているとわかる。

よって実装上は、ある dp[i+1][m] が2以上になったら更新を打ち切り、i+1から遡って集合Bを復元(要素a[i]を含む集合)、iから遡って集合Cを復元(要素a[i]を含まない集合)することになる。

上述した2通り以上になったら更新を打ち切ることはこの実装上の注意点となる。

数え上げDPをしていることになるので、Nが大きく、各要素の mod 200 の値がバラバラである場合などは、最後まで更新しようとするとLong型でもオーバーフローする可能性がある。(N=200なら、最悪2200に近い値を扱うことになる)

それを確認するには、以下のコードを実行してみると良い。

#include <bits/stdc++.h>
using namespace std;

#define rep(i, n) for (int i = (int)(0); i < (int)(n); ++i)
using ll = long long;
template<class t> using vc=vector<t>;

signed main()
{
  cin.tie( 0 ); ios::sync_with_stdio( false );
  
  vc<int> a(200);
  iota(a.begin(), a.end(),0);
  
  vc<vc<ll>> dp(201, vc<ll>(200));
  dp[0][0]=1;
  rep(i,200){
    rep(j,200){
      dp[i+1][j]+=dp[i][j];
      dp[i+1][(j+a[i])%200]+=dp[i][j];
    }
  }
}

runtime error: signed integer overflow: 5902957071859434052 + 5902959422909828636 cannot be represented in type 'long long int'

実装

#include <bits/stdc++.h>
using namespace std;

#define rep(i, n) for (int i = (int)(0); i < (int)(n); ++i)
#define irep(i, m, n) for (int i = (int)(m); i < (int)(n); ++i)
#define all(v) v.begin(), v.end()
#define sz(x) (int)(x.size())
template<class t> using vc=vector<t>;

signed main()
{
  cin.tie( 0 ); ios::sync_with_stdio( false );
  
  int n; cin>>n;
  vc<int> a(n);
  int all=0;
  rep(i,n){
    cin>>a[i];
    a[i]%=200;
    (all+=a[i])%=200;
  }
  
  bool allZero=0;
  if(all==0) allZero=1;
  
  vc<vc<int>> dp(n+1, vc<int>(200));
  dp[0][0]=1;
  rep(i,n){
    rep(j,200){
      int nj=(j+a[i])%200;
      dp[i+1][j]+=dp[i][j];
      dp[i+1][nj]+=dp[i][j];
    }
    if(dp[i+1][0]>1){
      if(i<n-1 or !allZero){
        // 部分集合で0を作れる
        vc<int> pZero; // 0(mod200)を作る部分集合
        set<int> s;
        int pre=0;
        for(int j=i; j>=0; j--){
          int tmp=(pre-a[j]+200)%200;
          if(dp[j][tmp]){
            pZero.push_back(j+1);
            s.insert(j);
            pre=tmp;
          }
        }
        int one=-1;
        rep(j,n) if(s.count(j)==0) {
          one=j+1;
          break;
        }
        
        cout<<"Yes"<<endl;
        cout<<1<<' '<<one<<endl;
        pZero.push_back(one);
        sort(all(pZero));
        cout<<pZero.size()<<' ';
        rep(j,pZero.size()) cout<<pZero[j]<<' ';
        cout<<endl;
        return 0;
      }
      
    }
    
    irep(j,1,200) if(dp[i+1][j]>1){
      // 部分集合0を使わずにB,Cを作れる
      vc<int> b,c;
      int pre=j;
      for(int k=i; k>=0; k--){
        int tmp=(pre-a[k]+200)%200;
        if(dp[k][tmp]){
          b.push_back(k+1);
          pre=tmp;
        }
      }
      pre=j;
      for(int k=i-1; k>=0; k--){
        int tmp=(pre-a[k]+200)%200;
        if(dp[k][tmp]){
          c.push_back(k+1);
          pre=tmp;
        }
      }
      reverse(all(b));
      reverse(all(c));
      cout<<"Yes"<<endl;
      cout<<b.size()<<' ';
      rep(k,b.size()) cout<<b[k]<<' ';
      cout<<endl;
      cout<<c.size()<<' ';
      rep(k,c.size()) cout<<c[k]<<' ';
      cout<<endl;
      return 0;
    }
  }
  
  cout<<"No"<<endl;
}

感想

400点問題にしては難しすぎると思ったけど、自身が難しく捉えているだけだった。