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

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

ACL Beginner Contest D - Flat Subsequence (水色, 400 点)

LIS を求める in-place DP を応用すればできる!
でも、400 点問題で「DP 配列をセグ木に乗せて」「in-place に更新することで高速化する」という問題が出るとは思わなかった!

in-place DP に馴染みのない方は先にこっちを

qiita.com

問題へのリンク

問題概要

数列  A_{1}, A_{2}, \dots, A_{N} と、正の整数  K が与えられる。

この数列の部分列であって、隣接する要素の差の絶対値が  K 以下であるようなものの、長さの最大値を求めよ。

制約

  •  N, K, A_{i} \le 3 \times 10^{5}

考えたこと

LIS を求めるアルゴリズムを知っていれば、それと同じようにできる。まずは素直な DP を考える。

  • dp[ i ][ v ] := 数列のうち最初の i 項のみを考えたときに、最後の項の値が v となる場合についての最長の長さ

このとき、a = A[ i ] とすると、

  • chmax( {\rm dp} \lbrack  i + 1 \rbrack \lbrack v \rbrack, {\rm dp} \lbrack  i  \rbrack \lbrack v \rbrack) (任意の  v)
  • chmax( {\rm dp} \lbrack  i + 1  \rbrack \lbrack a \rbrack, \max_{a - K \le j \le a + K} {\rm dp} \lbrack  i  \rbrack \lbrack j \rbrack)

という具合に遷移できる。そして前半の式を見ると次のことがわかる。


配列 dp[ i ] と、配列 dp[ i + 1 ] とで、同じ配列を使い回すようにすると、前者の chmax(dp[ i + 1 ][ v ], dp[ i ][ v ]) という更新をする必要はなくなる


つまり、後者の更新のみに注力すればよくなる。後者の更新において、 \max_{a - K \le j \le a + K} {\rm dp} \lbrack  i  \rbrack \lbrack j \rbrack の値を求めるのは、セグメント木を用いれば高速にできる。まとめると、

  • 配列 dp は、添字 i については共通化して使い回す
  • その配列 dp をセグメント木に乗せる

というテクニックによって  O(N \log N) の計算量で解くことができる。

#include <bits/stdc++.h>
using namespace std;
template<class T> inline bool chmax(T& a, T b) { if (a < b) { a = b; return 1; } return 0; }
template<class T> inline bool chmin(T& a, T b) { if (a > b) { a = b; return 1; } return 0; }

template<class Monoid> struct SegTree {
    using Func = function<Monoid(Monoid, Monoid)>;
    const Func F;
    const Monoid UNITY;
    int SIZE_R;
    vector<Monoid> dat;
    
    SegTree(int n, const Func f, const Monoid &unity): F(f), UNITY(unity) { init(n); }
    void init(int n) {
        SIZE_R = 1;
        while (SIZE_R < n) SIZE_R *= 2;
        dat.assign(SIZE_R * 2, UNITY);
    }
    
    /* set, a is 0-indexed */
    void set(int a, const Monoid &v) { dat[a + SIZE_R] = v; }
    void build() {
        for (int k = SIZE_R - 1; k > 0; --k)
            dat[k] = F(dat[k*2], dat[k*2+1]);
    }
    
    /* update a, a is 0-indexed */
    void update(int a, const Monoid &v) {
        int k = a + SIZE_R;
        dat[k] = v;
        while (k >>= 1) dat[k] = F(dat[k*2], dat[k*2+1]);
    }
    
    /* get [a, b), a and b are 0-indexed */
    Monoid get(int a, int b) {
        Monoid vleft = UNITY, vright = UNITY;
        for (int left = a + SIZE_R, right = b + SIZE_R; left < right; left >>= 1, right >>= 1) {
            if (left & 1) vleft = F(vleft, dat[left++]);
            if (right & 1) vright = F(dat[--right], vright);
        }                                                                                                              
        return F(vleft, vright);
    }
    inline Monoid operator [] (int a) { return dat[a + SIZE_R]; }
    
    /* debug */
    void print() {
        for (int i = 0; i < SIZE_R; ++i) {
            cout << (*this)[i];
            if (i != SIZE_R-1) cout << ",";
        }
        cout << endl;
    }
};

int solve(int N, int K) {
    const int MAX = 510000;
    auto func = [&](int a, int b) { return max(a, b); };
    SegTree<int> seg(MAX, func, 0);
    for (int i = 0; i < N; ++i) {
        int a;
        cin >> a;
        int left = max(0, a - K);
        int right = min(MAX, a + K + 1);
        int val = seg.get(left, right);
        seg.update(a, val + 1);
    }
    return seg.get(0, MAX);
}

int main() {
    int N, K;
    cin >> N >> K;
    cout << solve(N, K) << endl;
}