CODE FESTIVAL 2018 Final: G - Chicks and Cages

リンクを貼っていくスタイル。

atcoder.jp

小さいひよこほどケージに沢山詰め込むべき、というところまでは証明できた。だが、DP力が不足しておりそもそも dp(i,j) = j個のケージにi羽入れたときの最小値 という設定すら思いつかず、厳しい。

• 高速な言語だと、そもそも O(N^2 * M) 解が間に合う: http://kmjp.hatenablog.jp/entry/2018/12/06/0930
• O(NMlogN) 解の解説: http://fluffyowl.hatenablog.com/entry/2019/01/03/071824

上記解説を読んで実装した O(NMlogN) 解。if (size*j>k) break; という部分が本質。

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);

  int n, m; cin >> n >> m;
  vector<int> a(n);
  rep(i, n) {
    cin >> a[i];
  }
  sort(a.begin(), a.end());

  vector<vector<ll>> dp(n+1, vector<ll>(m+1, LONG_MAX/2));
  vector<ll> sum(n+1);
  for(int i=1; i<=n; i++) {
    sum[i] += sum[i-1] + a[i-1];
  }

  dp[0][0] = 0;
  for(int i=0; i<=n-1; i++) {
    for(int j=1; j<=m; j++) {
      for(int k=i+1; k<=n; k++) {
        auto size = k-i;
        if (size*j > k) break;
        chmin(dp[k][j], dp[i][j-1] + size*(sum[k] - sum[i]));
      }
    }
  }
  cout << dp[n][m] << endl;
}

添字を入れ替えると少し速くなる(45ms -> 38ms)

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);

  int n, m; cin >> n >> m;
  vector<int> a(n);
  rep(i, n) {
    cin >> a[i];
  }
  sort(a.begin(), a.end());

  vector<vector<ll>> dp(n+1, vector<ll>(m+1, LONG_MAX/2));
  vector<ll> sum(n+1);
  for(int i=1; i<=n; i++) {
    sum[i] += sum[i-1] + a[i-1];
  }

  dp[0][0] = 0;
  for(int i=0; i<=n-1; i++) {
    for(int k=i+1; k<=n; k++) {
      auto size = k-i;
      auto cost = (sum[k]-sum[i])*size;
      for(int j=1; j<=m; j++) {
        if (size*j > k) break;
        chmin(dp[k][j], dp[i][j-1]+cost);
      }
    }
  }
  cout << dp[n][m] << endl;
}