dp/1d1d.hpp

View this file on GitHub
Last update: 2024-11-09 00:49:07+00:00
Include: #include "dp/1d1d.hpp"

Depends on

prelude.hpp

Code

#pragma once
#include "prelude.hpp"

// O(n log n) (O(n) if we can skip binary search)
// Minimizes cost of partitions of [0, n)
// Requires cost(a,c) + cost(b,d) <= cost(a,d) + cost(b,c)
// for a < b < c < d (wider is worse).
template <class F>
auto dp_1d1d(int n, F cost) {
  using T = decltype(cost(0, 0));
  vector<T> dp(n+1);
  // opt[_] = argmin_{j<=k} (dp[j] + cost(j, i)) on k < ix[_] <= i < ix[_+1]
  vi opt(n+1), ix(n+1);
  int l = 0, r = ix[0]++; // opt[l:r+1] is a deque
  rep2(k, 1, n+1) {
    dp[k] = dp[opt[l]] + cost(opt[l], k);
    assert(l+1 < n);
    if (++ix[l] == ix[l+1] && l < r) l++;
    while (l <= r && dp[opt[r]] + cost(opt[r], ix[r]) > dp[k] + cost(k, ix[r])) r--;
    if (opt[++r] = k, l == r) continue;
    int lo = ix[--r], hi = n+1; // Finds hi = where opt's value becomes k
    while (hi - lo > 1) {
      int m = (lo + hi) / 2;
      (dp[opt[r]] + cost(opt[r], m) > dp[k] + cost(k, m) ? hi : lo) = m;
    }
    if (hi <= n) opt[++r] = k, ix[r] = hi;
  }
  return dp;
}

#line 2 "prelude.hpp"
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using vi = vector<int>;
using vvi = vector<vector<int>>;
using vll = vector<ll>;
using vvll = vector<vector<ll>>;
using vc = vector<char>;
#define rep2(i, m, n) for (auto i = (m); i < (n); i++)
#define rep(i, n) rep2(i, 0, n)
#define repr2(i, m, n) for (auto i = (n); i-- > (m);)
#define repr(i, n) repr2(i, 0, n)
#define all(x) begin(x), end(x)
auto ndvec(int n, auto e) { return vector(n, e); }
auto ndvec(int n, auto ...e) { return vector(n, ndvec(e...)); }
auto comp_key(auto&& f) { return [&](auto&& a, auto&& b) { return f(a) < f(b); }; }
auto& max(const auto& a, const auto& b) { return a < b ? b : a; }
auto& min(const auto& a, const auto& b) { return b < a ? b : a; }
#if __cpp_lib_ranges
namespace R = std::ranges;
namespace V = std::views;
#endif
#line 3 "dp/1d1d.hpp"

// O(n log n) (O(n) if we can skip binary search)
// Minimizes cost of partitions of [0, n)
// Requires cost(a,c) + cost(b,d) <= cost(a,d) + cost(b,c)
// for a < b < c < d (wider is worse).
template <class F>
auto dp_1d1d(int n, F cost) {
  using T = decltype(cost(0, 0));
  vector<T> dp(n+1);
  // opt[_] = argmin_{j<=k} (dp[j] + cost(j, i)) on k < ix[_] <= i < ix[_+1]
  vi opt(n+1), ix(n+1);
  int l = 0, r = ix[0]++; // opt[l:r+1] is a deque
  rep2(k, 1, n+1) {
    dp[k] = dp[opt[l]] + cost(opt[l], k);
    assert(l+1 < n);
    if (++ix[l] == ix[l+1] && l < r) l++;
    while (l <= r && dp[opt[r]] + cost(opt[r], ix[r]) > dp[k] + cost(k, ix[r])) r--;
    if (opt[++r] = k, l == r) continue;
    int lo = ix[--r], hi = n+1; // Finds hi = where opt's value becomes k
    while (hi - lo > 1) {
      int m = (lo + hi) / 2;
      (dp[opt[r]] + cost(opt[r], m) > dp[k] + cost(k, m) ? hi : lo) = m;
    }
    if (hi <= n) opt[++r] = k, ix[r] = hi;
  }
  return dp;
}