5.3.2 Divide-and-Conquer DP Optimization

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89

/*

Computes one layer of a dynamic programming recurrence whose optimal transition indices are
monotone. This optimization applies to recurrences of the form
`dp_cur[i] = min(dp_prev[k] + cost(k, i))`, where the minimizing index for $i$ never decreases as
$i$ increases.

The function below computes `dp_cur[l..r]` by evaluating the midpoint first, then recursing on the
left half with a restricted upper bound on the optimal `k`, and on the right half with a restricted
lower bound. The caller is responsible for checking that the recurrence satisfies the required
monotonicity condition.

- `compute_dc_layer(dp_prev, dp_cur, l, r, opt_l, opt_r, cost)` fills `dp_cur[l..r]` using candidate
  transition indices in [`opt_l`, `opt_r`]. The template parameter `cost` must be callable such that
  `cost(k, i)` returns the transition cost from previous state `k` to current state `i`.

Time Complexity:
- O(n log n) calls to `cost(k, i)` per layer when each state has O(n) possible transitions.

Space Complexity:
- O(log n) auxiliary stack space.

*/

#include <algorithm>
#include <cstdint>
#include <vector>

const int64_t INF = (1LL << 62);

template<class Cost>
void compute_dc_layer(
    const std::vector<int64_t> &dp_prev, std::vector<int64_t> &dp_cur, int l, int r, int opt_l,
    int opt_r, Cost cost
) {
  if (l > r) {
    return;
  }
  int mid = l + (r - l) / 2;
  int64_t best = INF;
  int best_k = opt_l;
  int upper = std::min(mid, opt_r);
  for (int k = opt_l; k <= upper; k++) {
    int64_t candidate = dp_prev[k] + cost(k, mid);
    if (candidate < best) {
      best = candidate;
      best_k = k;
    }
  }
  dp_cur[mid] = best;
  compute_dc_layer(dp_prev, dp_cur, l, mid - 1, opt_l, best_k, cost);
  compute_dc_layer(dp_prev, dp_cur, mid + 1, r, best_k, opt_r, cost);
}

/*** Example Usage ***/

#include <cassert>
using namespace std;

struct SquareSegmentCost {
  vector<int64_t> prefix;

  explicit SquareSegmentCost(const vector<int> &a) : prefix(a.size() + 1) {
    for (int i = 0; i < static_cast<int>(a.size()); i++) {
      prefix[i + 1] = prefix[i] + a[i];
    }
  }

  int64_t operator()(int k, int i) const {
    int64_t sum = prefix[i] - prefix[k];
    return sum * sum;
  }
};

int main() {
  vector<int> a{1, 2, 3, 4};
  int n = static_cast<int>(a.size());
  SquareSegmentCost cost(a);

  vector<int64_t> dp_prev(n + 1, INF), dp_cur(n + 1, INF);
  dp_prev[0] = 0;
  compute_dc_layer(dp_prev, dp_cur, 1, n, 0, n - 1, cost);
  assert(dp_cur[4] == 100);  // One group: (1 + 2 + 3 + 4)^2.

  vector<int64_t> dp_two(n + 1, INF);
  compute_dc_layer(dp_cur, dp_two, 1, n, 0, n - 1, cost);
  assert(dp_two[4] == 52);  // Split as [1, 2] and [3, 4].
  return 0;
}