JOI 2015/2016 本選 A - オレンジの出荷 (Oranges)

JOI 2015/2016 本選 A - オレンジの出荷 (Oranges) に関連する、再帰呼び出し・メモ化再帰、非再帰についてのメモを残します。

とにかく再帰が苦手

• 問題 : JOI 2015/2016 本選 A - オレンジの出荷 (Oranges)

入力例 1 : 、、 について考える。
1 番のオレンジからみていくが、箱には 3 個まで詰めることができるため、1 個目の箱の詰め方は、

• 
• 
• 

以上の 3 通り。それぞれについて、再帰的に 2 個目以降の詰め方を全探索する。

まずはメモ化を行わずに書いてみる。
メモ化をしないと、上図のように同じ計算を何回も行う無駄が生じてしまい、今回の問題では TLE になる。

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

int N, M, K;
vector<long long> A;
const long long INF = 1LL << 60;

long long func(int idx) {
  if (idx == N) return 0;  // ベースケース
  long long res = INF;
  long long max_a = -INF, min_a = INF;
  // 箱に何個詰めるか
  for (int i = idx; i < idx + M; i++) {
    if (i == N) break;
    max_a = max(max_a, A.at(i));
    min_a = min(min_a, A.at(i));
    long long tmp = K + (i - idx + 1) * (max_a - min_a);
    long long cost = tmp + func(i+1);
    res = min(res, cost);
  }
  return res;
}

int main() {
  cin >> N >> M >> K;
  A.resize(N);
  for (int i = 0; i < N; i++) cin >> A.at(i);
  cout << func(0) << endl;
}

比較として、途中までのコストを引数にもつバージョンを下に載せておく。
ただし、これだとうまくメモ化することができない。

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

int N, M, K;
vector<long long> A;
const long long INF = 1LL << 60;

long long func(int idx, long long cost) {
  if (idx == N) return cost;
  long long res = INF;
  long long max_a = -INF, min_a = INF;
  for (int i = idx; i < idx + M; i++) {
    if (i == N) break;
    max_a = max(max_a, A.at(i));
    min_a = min(min_a, A.at(i));
    long long tmp = K + (i - idx + 1) * (max_a - min_a);
    res = min(res, func(i+1, cost + tmp));
  }
  return res;
}

int main() {
  cin >> N >> M >> K;
  A.resize(N);
  for (int i = 0; i < N; i++) cin >> A.at(i);
  cout << func(0, 0) << endl;
}

メモ化再帰

変わる部分は少しだけ。

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

int N, M, K;
vector<long long> A;
vector<long long> dp;
const long long INF = 1LL << 60;

long long func(int idx) {
  if (idx == N) return 0;  // ベースケース
  if (dp.at(idx) != INF) return dp.at(idx);  // メモ済み
  long long res = INF;
  long long max_a = -INF, min_a = INF;
  // 箱に何個詰めるか
  for (int i = idx; i < idx + M; i++) {
    if (i == N) break;
    max_a = max(max_a, A.at(i));
    min_a = min(min_a, A.at(i));
    long long tmp = K + (i - idx + 1) * (max_a - min_a);
    long long cost = tmp + func(i+1);
    res = min(res, cost);
  }
  return dp.at(idx) = res;  // メモ
}

int main() {
  cin >> N >> M >> K;
  A.resize(N);
  dp.assign(N, INF);
  for (int i = 0; i < N; i++) cin >> A.at(i);
  cout << func(0) << endl;
}

非再帰で実装

メモ化再帰のコードから非再帰のコードを書いてみた。

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

int main() {
  int N, M, K; cin >> N >> M >> K;
  vector<long long> A(N);
  for (int i = 0; i < N; i++) cin >> A.at(i);
  const long long INF = 1LL << 60;
  // dp[i] := i 番目以降のオレンジを詰めるときの最小コスト
  vector<long long> dp(N+1, INF);
  dp.at(N) = 0;
  for (int i = N; i > 0; i--) {
    long long max_a = -INF, min_a = INF;
    for (int j = 1; j <= M; j++) {
      if (i-j < 0) break;
      max_a = max(max_a, A.at(i-j));
      min_a = min(min_a, A.at(i-j));
      long long cost = dp.at(i) + K + j * (max_a - min_a);
      dp.at(i-j) = min(dp.at(i-j), cost);
    }
  }
  cout << dp.at(0) << endl;
}