test/algorithm/monge_2.test.cpp

View this file on GitHub
Last update: 2026-10-21 23:29:36+09:00
Problem: https://yukicoder.me/problems/no/952

Depends on

Code

#define PROBLEM "https://yukicoder.me/problems/no/952"

#include "../../algorithm/monge_k_edge_shortest_path.hpp"
#include <iostream>
#include <vector>

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

    int N;
    std::cin >> N;
    std::vector<long long> S(N + 1);
    for (int i = 0; i < N; ++i){
        std::cin >> S[i + 1];
        S[i + 1] += S[i];
    }

    auto cost = [&](int i, int j){
        long long sum = S[j - 1] - S[i];
        return sum * sum;
    };

    auto ans = po167::monge_k_edge_shortest_path<long long>(N + 1, N, cost);
    for (int k = 1; k <= N; ++k){
        std::cout << ans[N - k + 1] << "\n";
    }
}

#line 1 "test/algorithm/monge_2.test.cpp"
#define PROBLEM "https://yukicoder.me/problems/no/952"

#line 2 "algorithm/monge_k_edge_shortest_path.hpp"
// po167 Monge library v2 (2026-10-21)
#include <cassert>
#include <vector>
#line 3 "algorithm/smawk.hpp"
#include <numeric>
namespace po167{
// return argmax
// 
// https://noshi91.github.io/Library/algorithm/smawk.cpp
template<class S> std::vector<int> smawk(int N, int M, S select){
    auto f = [&](auto self, std::vector<int> p,std::vector<int> q) -> std::vector<int> {
        // reduce
        std::vector<int> n_q;
        n_q.reserve(p.size());
        int k = -1;
        for(auto x:q){
            while(!n_q.empty()){
                if(!select(p[k], n_q.back(), x)) break;
                n_q.pop_back();
                k--;
            }
            if (k + 1 != (int)p.size()){
                n_q.push_back(x);
                k++;
            }
        }
        std::swap(n_q, q);

        int n = p.size();
        int m = q.size();
        
        // recursive
        if (n == 1){
            return q;
        }
        std::vector<int> n_p;
        n_p.reserve(n / 2);

        for (int i = 1; i < n; i += 2){
            n_p.push_back(p[i]);
        }
        
        auto fr = self(self, n_p, q);
        // interpolate
        std::vector<int> ans(n);
        int l = 0, r = 0;
        
        for (int i = 0; i < n; i++){
            if(i & 1){
                ans[i] = fr[i / 2];
                l = r;
            } else {
                if (i + 1 != n) {
                    while (q[r] != fr[i / 2]) r++;
                } else {
                    r = m - 1;
                }
                ans[i] = q[l];
                for (int j = l + 1; j <= r; j++){
                    if(select(p[i], ans[i], q[j])){
                        ans[i] = q[j];
                    }
                }
            }
        }
        return ans;
    };
    std::vector<int> P(N), Q(M);
    std::iota(P.begin(), P.end(), 0);
    std::iota(Q.begin(), Q.end(), 0);
    return f(f, P, Q);
}
}
#line 6 "algorithm/monge_k_edge_shortest_path.hpp"

namespace po167 {

// f: {(i,j) | 0 <= i < j <= n} -> long long, 0 <= k <= n.
// f(a,c) + f(b,d) <= f(a,d) + f(b,c)  (a < b < c < d).
// ans[d] = min_{0 = v_0 < ... < v_d = n} sum_{t=1}^d f(v_{t-1},v_t), 1 <= d <= k.
// ans[0] は未使用。各 ans[d] は long long、全 DP 値・候補値は T に収まること。
// f が O(1) のとき、時間 O(nk)、追加メモリ O(n)。
template <class T = __int128_t, class F>
std::vector<long long> monge_k_edge_shortest_path(int n, int k, const F& f) {
    assert(0 <= k && k <= n);
    std::vector<long long> answer(k + 1);
    if (k == 0) return answer;
    answer[1] = f(0, n);
    if (k == 1) return answer;

    std::vector<T> prev(n + 1), next(n + 1);
    for (int j = 1; j < n; ++j) prev[j] = T(f(0, j));
    prev[n] = T(answer[1]);
    for (int d = 2; d <= k; ++d) {
        int size = n - d + 1;
        auto argmin = po167::smawk(size, size, [&](int r, int a, int b) {
            int j = r + d, i = a + d - 1, h = b + d - 1;
            return h < j && prev[h] + T(f(h, j)) < prev[i] + T(f(i, j));
        });
        for (int j = d; j <= n; ++j) {
            int i = argmin[j - d] + d - 1;
            next[j] = prev[i] + T(f(i, j));
        }
        answer[d] = (long long)next[n];
        prev.swap(next);
    }
    return answer;
}

}  // namespace po167
#line 4 "test/algorithm/monge_2.test.cpp"
#include <iostream>
#line 6 "test/algorithm/monge_2.test.cpp"

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

    int N;
    std::cin >> N;
    std::vector<long long> S(N + 1);
    for (int i = 0; i < N; ++i){
        std::cin >> S[i + 1];
        S[i + 1] += S[i];
    }

    auto cost = [&](int i, int j){
        long long sum = S[j - 1] - S[i];
        return sum * sum;
    };

    auto ans = po167::monge_k_edge_shortest_path<long long>(N + 1, N, cost);
    for (int k = 1; k <= N; ++k){
        std::cout << ans[N - k + 1] << "\n";
    }
}