けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

COLOCON 2018 Final C - スペースエクスプローラー高橋君 (600 点)

Convex Hull Trick の練習に。DP ではないけど、DP 高速化でよくやる Convex Hull Trick の構造そのものを試せる問題。Convex Hull Trick についての詳しい解説は

がとてもわかりやすい!

問題へのリンク

問題概要

長さ  N の数列  a_1, a_2, \dots, a_N が与えられる。各  i に対して、

  •  {\rm min}_{j} (a_{j} + (i - j)^{2})

の値を求めよ。

制約

  •  1 \le N \le 2×10^{5}

解法 1: Convex Hull Trick (直線の傾きも、クエリ値にも単調性がある)

Convex Hull Trick

  • insert (a, b): 直線  y = ax + b を追加する
  • query (x):  {\rm min}_{i} (a_{i}*x + b_{i}) を答える

を高速にできるデータ構造で、一般にともに  O(\log{n}) でできる。しかし嬉しい仮定として

  • insert される順番に  a が単調減少
  • query される順番に  x が単調増加

の場合はならし  O(1) でできる (公式解説参照)。

今回の問題で言えば、

 {\rm min}_{j} (a_{j} + (i - j)^{2}) = {\rm min}_{j} { (-2j)i + (a_{j} + j^{2})} +  i^{2}

と式変形すると、

  • 直線集合は各  j に対して  y = (-2j) x + (a_{j} + j^{2})
  • クエリは  x = i を各  i に対して

という感じになる。全体の計算量は  O(N)

#include <iostream>
#include <vector>
#include <deque>
#include <algorithm>
using namespace std;


// Convex Hull Trick
/*
 追加クエリの直線の傾きが単調減少の場合
 
 - insert (a, b): add y = ax + b
 - query (x): min_i{a[i]*x + b}
 */
template<class T> struct CHT {
    using Line = pair<T, T>;
    deque<Line> lines;
    
    inline bool check(const Line &a, const Line &b, const Line &c) {
        return (b.first - a.first) * (c.second - b.second)
        >= (b.second - a.second) * (c.first - b.first);
    }
    
    void insert(T a, T b) {
        Line l(a, b);
        while (lines.size() >= 2
               && check(lines[(int)lines.size()-2], lines[(int)lines.size()-1], l))
            lines.pop_back();
        lines.push_back(l);
    }
    
    T query(T x) {
        
    }
    
    // クエリの単調性も成り立つ場合 (x が単調増加)
    T query_monotone(T x) {
        while (lines.size() >= 2
               && lines[0].first * x + lines[0].second >= lines[1].first * x + lines[1].second)
            lines.pop_front();
        return lines[0].first * x + lines[0].second;
    }
};

int main() {
    long long N; cin >> N;
    vector<long long> a(N), res(N, 1LL<<60);
    for (int i = 0; i < N; ++i) cin >> a[i];
    CHT<long long> cht;
    for (long long i = 0; i < N; ++i) cht.insert(-2LL*i, a[i] + i*i);
    for (long long i = 0; i < N; ++i) res[i] = cht.query_monotone(i) + i*i;
    for (int i = 0; i < N; ++i) cout << res[i] << endl;
}

解法 2: Convex Hull Trick (直線の傾きに単調性がある)

直線の傾きに単調性があることを利用した実装をする。query を投げるときに二分探索をする。詳しくは、さたにゃんの解説記事を読むとわかりやすい。

全体の計算量は  O(N\log{N})

#include <iostream>
#include <vector>
#include <deque>
#include <algorithm>
using namespace std;


// Convex Hull Trick
/*
    追加クエリの直線の傾きが単調減少の場合
    
    - insert (a, b): add y = ax + b
    - query (x): min_i{a[i]*x + b}
*/
template<class T> struct CHT {
    using Line = pair<T, T>;
    deque<Line> lines;
    
    inline bool check(const Line &a, const Line &b, const Line &c) {
        return (b.first - a.first) * (c.second - b.second)
        >= (b.second - a.second) * (c.first - b.first);
    }
    
    inline T get(int k, T x) {
        return lines[k].first * x + lines[k].second;
    }
    
    void insert(T a, T b) {
        Line l(a, b);
        while (lines.size() >= 2
               && check(lines[(int)lines.size()-2], lines[(int)lines.size()-1], l))
            lines.pop_back();
        lines.push_back(l);
    }
    
    T query(T x) {
        int low = -1, high = (int)lines.size();
        while (high - low > 1) {
            int mid = (low + high) / 2;
            if (get(mid, x) >= get(mid + 1, x)) low = mid;
            else high = mid;
        }
        return get(high, x);
    }
    
    // クエリの単調性も成り立つ場合 (x が単調増加)
    T query_monotone(T x) {
        while (lines.size() >= 2 && get(0, x) >= get(1, x)) lines.pop_front();
        return lines[0].first * x + lines[0].second;
    }
};

int main() {
    long long N; cin >> N;
    vector<long long> a(N), res(N, 1LL<<60);
    for (int i = 0; i < N; ++i) cin >> a[i];
    CHT<long long> cht;
    for (long long i = 0; i < N; ++i) cht.insert(-2LL*i, a[i] + i*i);
    for (long long i = 0; i < N; ++i) res[i] = cht.query(i) + i*i;
    for (int i = 0; i < N; ++i) cout << res[i] << endl;
}

解法 3: Convex Hull Trick (直線の単調性もない)

Li Chao Segment Tree を用いることで直線の単調性が成り立たない場合にも使えるようにする。以下の記事を参考に:

kazuma さんの記事ではクエリ先読みに対応する形で書いてあるが、 x の範囲を先に定めることでオンラインクエリで対応できることが、てんぷらさん・ferin さんによって示された。計算量は  O(N\log{N})

#include <iostream>
#include <vector>
#include <deque>
#include <algorithm>
using namespace std;


// Convex Hull Trick
/*
    - insert (a, b): add y = ax + b
    - query (x): min_i{a[i]*x + b}
*/
template<class T> struct CHT {
    struct Line {
        T a, b;
        Line(T a = 0,T b = 0) : a(a), b(b) { }
        T get(T x) {
            return a*x + b;
        }
    };
    
    struct Node {
        Line l;
        Node *lhs,*rhs;
        Node(Line l) : l(l), lhs(nullptr), rhs(nullptr){}
    };
    
    const T MIN, MAX, INF;
    Node* root;
    
    CHT(T MIN, T MAX, T INF) : MIN(MIN), MAX(MAX), INF(INF), root(nullptr) { }
    Node* insert(Node* p, T left, T right, Line& l){
        if(!p) return new Node(l);
        if (p->l.get(left) <= l.get(left) && p->l.get(right) <= l.get(right)) return p;
        if (p->l.get(left) >= l.get(left) && p->l.get(right) >= l.get(right)) {
            p->l = l;
            return p;
        }
        T mid = (left + right) / 2;
        if (p->l.get(mid) > l.get(mid)) swap(p->l , l);
        if (p->l.get(left) >= l.get(left)) p->lhs = insert(p->lhs, left, mid, l);
        else p->rhs = insert(p->rhs, mid+1, right, l);
        return p;
        
    }
    void insert(T a, T b){
        Line l(a, b);
        root = insert(root, MIN, MAX, l);
    }
    T query(Node* p, T left, T right, T x){
        if(!p) return INF;
        if(left == right) return p->l.get(x);
        T mid = (left + right) / 2;
        if (x <= mid) return min(p->l.get(x), query(p->lhs, left, mid, x));
        else return min(p->l.get(x), query(p->rhs, mid+1, right, x));
    }
    T query(T x){
        return query(root, MIN, MAX, x);
    }
};


int main() {
    long long N; cin >> N;
    vector<long long> a(N), res(N);
    for (int i = 0; i < N; ++i) cin >> a[i];
    CHT<long long> cht(0, 1LL<<30, 1LL<<60);
    for (long long i = 0; i < N; ++i) cht.insert(-2LL*i, a[i] + i*i);
    for (long long i = 0; i < N; ++i) res[i] = cht.query(i) + i*i;
    for (int i = 0; i < N; ++i) cout << res[i] << endl;
}

解法 4: 一般化 CHT

よくよく考えてみると、CHT に追加するものは直線でなくてもよくて

  • どの 2 つをとっても、共有点は高々 1 個である

という性質を満たしていればよい。したがって、通常の CHT では一次式の最小値になるように頑張るところを、直接二次関数をぶっこんでしまってもよい。CHT ライブラリの node->get の中身を書き換える。

#include <iostream>
#include <vector>
#include <deque>
#include <algorithm>
using namespace std;


// Convex Hull Trick
template<class T> struct CHT {
    struct Line {
        T a, b;
        Line(T a = 0,T b = 0) : a(a), b(b) { }
        T get(T x) {
            return a + (x - b) * (x - b);
        }
    };
    
    struct Node {
        Line line;
        Node *left, *right;
        Node(Line line) : line(line), left(nullptr), right(nullptr){}
    };
    
    const T MIN, MAX, INF;
    Node* root;
    
    CHT(T MIN, T MAX, T INF) : MIN(MIN), MAX(MAX), INF(INF), root(nullptr) { }
    Node* insert(Node* p, T low, T high, Line& l){
        if(!p) return new Node(l);
        if (p->line.get(low) <= l.get(low) && p->line.get(high) <= l.get(high)) return p;
        if (p->line.get(low) >= l.get(low) && p->line.get(high) >= l.get(high)) {
            p->line = l;
            return p;
        }
        T mid = (low + high) / 2;
        if (p->line.get(mid) > l.get(mid)) swap(p->line , l);
        if (p->line.get(low) >= l.get(low)) p->left = insert(p->left, low, mid, l);
        else p->right = insert(p->right, mid, high, l);
        return p;
        
    }
    void insert(T a, T b){
        Line l(a, b);
        root = insert(root, MIN, MAX, l);
    }
    T query(Node* p, T low, T high, T x){
        if(!p) return INF;
        if(low == high) return p->line.get(x);
        T mid = (low + high) / 2;
        if (x <= mid) return min(p->line.get(x), query(p->left, low, mid, x));
        else return min(p->line.get(x), query(p->right, mid, high, x));
    }
    T query(T x){
        return query(root, MIN, MAX, x);
    }
};

int main() {
    long long N; cin >> N;
    vector<long long> a(N), res(N, 1LL<<60);
    for (int i = 0; i < N; ++i) cin >> a[i];
    CHT<long long> cht(0, N+1, 1LL<<60);
    for (long long i = 0; i < N; ++i) cht.insert(a[i], i);
    for (long long i = 0; i < N; ++i) res[i] = cht.query(i);
    for (int i = 0; i < N; ++i) cout << res[i] << endl;
}