5.3.1 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
90
91
92
93
94
95
96
97
98

/*

Computes one layer $t$ of a dynamic program of the form
$dp_t(i) = \min_{0 \leq k \leq i} (dp_{t-1}(k) + C(k, i))$, where $C(k, i)$ is the cost of moving
from state $k$ of the previous layer to state $i$ of the current one. Let $opt(i)$ be the smallest
$k$ attaining this minimum. Divide-and-conquer optimization applies when $opt(i) \leq opt(i + 1)$,
so the best transition index moves only to the right as $i$ increases.

To compute $dp_t$ over a range of states, evaluate its midpoint and record the best transition index
$k^*$ found there. Monotonicity restricts every optimum in the left half to indices at most $k^*$,
and every optimum in the right half to indices at least $k^*$. Recursively applying these bounds
avoids checking every transition for every state. The caller must verify the required monotonicity
property.

- `compute_dp_layer(dp_prev, dp_cur, lo, hi, opt_lo, opt_hi, cost)` fills `dp_cur[lo..hi]` using
  candidate transition indices in [`opt_lo`, `opt_hi`], where `dp_prev` holds $dp_{t-1}$ and
  `dp_cur` receives $dp_t$. The template parameter `cost` must be callable such that `cost(k, i)`
  returns $C(k, i)$, the transition cost from previous state `k` to current state `i`.

Time Complexity:
- O(n log n) calls to `cost()` per call, when computing one layer whose states each have O(n)
  possible transitions.

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

*/

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

const int64_t INF = INT64_MAX / 4;

template<typename Cost>
void compute_dp_layer(
    const std::vector<int64_t> &dp_prev, std::vector<int64_t> &dp_cur, int lo, int hi, int opt_lo,
    int opt_hi, const Cost &cost
) {
  if (lo > hi) {
    return;
  }
  int mid = lo + (hi - lo) / 2;
  int64_t best = INF;
  int best_k = opt_lo;
  int upper = std::min(mid, opt_hi);
  for (int k = opt_lo; k <= upper; k++) {
    int64_t candidate = dp_prev[k] + cost(k, mid);  // Overflow warning.
    if (candidate < best) {
      best = candidate;
      best_k = k;
    }
  }
  dp_cur[mid] = best;
  compute_dp_layer(dp_prev, dp_cur, lo, mid - 1, opt_lo, best_k, cost);
  compute_dp_layer(dp_prev, dp_cur, mid + 1, hi, best_k, opt_hi, 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 {
    // Cost of placing the half-open segment a[k..i) in one group.
    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);

  // Base layer for zero groups: only the empty prefix is reachable.
  vector<int64_t> dp_prev(n + 1, INF), dp_cur(n + 1, INF);
  dp_prev[0] = 0;

  // The first layer places each prefix a[0..j) into one group.
  compute_dp_layer(dp_prev, dp_cur, 1, n, 0, n - 1, cost);
  assert(dp_cur[4] == 100);  // One group: (1 + 2 + 3 + 4)^2.

  // The second layer chooses the best split between two groups.
  vector<int64_t> dp_two(n + 1, INF);
  compute_dp_layer(dp_cur, dp_two, 1, n, 0, n - 1, cost);
  assert(dp_two[4] == 52);  // (1 + 2 + 3)^2 + 4^2.
  return 0;
}