Official

B - At Most 3 (Judge ver.) Editorial by Nyaan


まず、この問題文をプログラミング的な手続きに言い換えると次のような問題になります。この手順を十分高速に計算できればこの問題を解くことができます。

  • \(1\) 以上 \(W\) 以下のすべての整数について「\(n\) はよい整数か」を管理するフラグ用の配列 flag を用意する。flag ははじめ false で初期化されている。
  • \(3\) 個以下であるおもりの集合全てを調べる。各集合についておもりの重さの和 \(w\) を求める。そして、\(w\)\(W\) 以下ならば flag[w]true にする。
  • 最終的に true であるフラグの個数が答えとなる。

上の手順の中で一番難しいのが「 \(3\) 個以下であるおもりの集合全てを調べる」という部分で、ここで探索の方法を間違えると計算量が膨大になって TLE しています。
この部分は「\(3\) 個以下である」を「\(1\) 個 または \(2\) 個または \(3\) 個である」と言い換えるのがポイントです。「おもりの個数がちょうど \(x\) 個である集合」の場合は 以下のような for-loop を用いた実装で \(\mathrm{O}(N^x)\) で計算することができるので、このように言い換えたあと \(1\) 個の場合, \(2\) 個の場合, \(3\) 個の場合を別々に計算すれば、この問題は for-loop で実装できる問題に帰着します。

// for-loop の x = 2 のときの例
  for(int i = 0; i < N; i++) {
    for(int j = i + 1; j < N; j++) {
      // 異なるおもりの組 (i, j) について計算する
    }
  }

計算量はフラグ用配列の分の \(\mathrm{O}(W)\) と for-loop の \(\mathrm{O}(N+N^2+N^3) = \mathrm{O}(N^3)\) を合わせて \(\mathrm{O}(W + N^3)\) となり、これは十分高速に動作します。

C++ による実装例は次の通りです。

#include <iostream>
#include <vector>
using namespace std;

int main() {
  int N, W;
  cin >> N >> W;
  vector<int> A(N);
  for (auto& x : A) cin >> x;
  
  vector<int> flag(W + 1);

  for (int i = 0; i < N; i++) {
    int s = A[i];
    if (s <= W) flag[s] = 1;
  }

  for (int i = 0; i < N; i++) {
    for (int j = i + 1; j < N; j++) {
      int s = A[i] + A[j];
      if (s <= W) flag[s] = 1;
    }
  }
  
  for (int i = 0; i < N; i++) {
    for (int j = i + 1; j < N; j++) {
      for (int k = j + 1; k < N; k++) {
        int s = A[i] + A[j] + A[k];
        if (s <= W) flag[s] = 1;
      }
    }
  }
  
  int ans = 0;
  for (int i = 1; i <= W; i++) ans += flag[i];
  cout << ans << endl;
}

posted:
last update: