Official

N - ジグザグな数列/Zigzag Sequences Editorial by penguinman


\(A\) の先頭から使う要素を選んでいき、ジグザグな数列になるような部分列 \(B\) を構築していくことを考えます。

手始めに、以下のような動的計画法を考えてみましょう。便宜上 \(2 \leq |B|\) を仮定した定義となっていますが、\(|B|=1\) であるような場合は別途適切に処理してください。

  • \(\text{dp}[i][j]:=(B\) の末尾の要素が \(A_i\) であり、かつ \(j=0\) なら \((B\) の末尾から \(2\) 番目の要素が \(x\) であるとして\()\) \(A_i \lt x\) が、\(j=1\) なら \(x \lt A_i\) が満たされているような \(B\) の通り数\()\)

上記の動的計画法は以下のような漸化式で表されます。

  • \(\text{dp}[i][0]:=\sum_{j \lt i, A_i \lt A_j} \text{dp}[j][1]\)
  • \(\text{dp}[i][1]:=\sum_{j \lt i, A_j \lt A_i} \text{dp}[j][0]\)

このような動的計画法は愚直に行うと \(O(N^2)\) 回の計算を要しますが、数列 \(A\) を座標圧縮した上で Segment Tree を利用すると計算回数を \(O(N \log N)\) に抑えることができます。

よってこの問題を解くことができました。

実装例 (C++)

#include<bits/stdc++.h>
using std::cin;
using std::cout;
using std::endl;
using std::vector;
using std::string;
using ll = long long;
#define rep(i,j,k) for(int i=int(j); i<int(k); i++)
#define REP(i,j,k) for(int i=int(j); i<=int(k); i++)
#define per(i,j,k) for(int i=int(j); i>=int(k); i--)
const ll mod = 1e9+7;

//Segment Tree
template<typename T>
struct segment_tree{
    ll N;
    T INF;
    vector<T> node;
    segment_tree(vector<T> X, T INF_):INF(INF_){
        N = 1;
        while(N < X.size()) N <<= 1;
        node.resize(2*N, INF);
        rep(i, 0, X.size()) node[N+i] = X[i];
        per(i, N-1, 0) node[i] = compare(node[i*2], node[i*2+1]);
    }
    T compare(T A, T B){
        return (A+B)%mod;
    }
    //isAdd == true: add
    //isAdd == false: update
    void update(ll index, T val, bool isAdd){
        index += N;
        if(isAdd) node[index] += val;
        else node[index] = val;
        while(index > 0){
            index >>= 1;
            node[index] = compare(node[index*2], node[index*2+1]);
        }
    }
    T calc(ll left, ll right){
        T ret = INF;
        left = std::max(0ll, left);
        right = std::min(N, right);
        left += N;
        right += N;
        while(left < right){
            if(left & 1) ret = compare(node[left++], ret);
            if(right & 1) ret = compare(node[--right], ret);
            left >>= 1;
            right >>= 1;
        }
        return ret;
    }
};

//座標圧縮
vector<int> comp(vector<int> A){
    std::map<int,int> mem;
    for(auto p: A) mem[p] = 0;
    int cnt = 0;
    for(auto &p: mem) p.second = cnt++;
    vector<int> B = A;
    for(auto &p: B) p = mem[p];
    return B;
}

int main(){
    int N; cin >> N;
    vector<int> A(N);
    rep(i,0,N) cin >> A[i];
    A = comp(A);
    ll ans = 0;
    vector<segment_tree<ll>> seg = {segment_tree<ll>(vector<ll>(N),0),segment_tree<ll>(vector<ll>(N),0)};
    for(auto p: A){
        ll a = seg[0].calc(0,p);
        ll b = seg[1].calc(p+1,N);
        ans += a+b;
        ans %= mod;
        a++;
        a %= mod;
        b++;
        b %= mod;
        seg[0].update(p,b,true);
        seg[1].update(p,a,true);
    }
    cout << ans << endl;
}

posted:
last update: