Codeforces Round #361 (Div. 2) D. Friends and Subsequences

問題文

http://codeforces.com/contest/689/problem/D

問題概要

長さ N の数列 A, B が与えられる．
max{A(l), ..., A(r)} = min{B(l), ... B(r)} となる (l, r) (l <= r) は何組存在するか求めよ．

・制約
1 <= N <= 200000

解法

Sparse table を書いてみたのでそのテストに使った．なので Sparse table で解く．
まず適当に l を定めて，条件を満たす右端の範囲を二分探索で求める．
l を定めたときに，l を左端として条件を満たすものが存在する可能性がなければ continue する．
存在しうるなら
l2 で初めて max{A(l), ..., A(l2)} >= min{B(l), ..., B(l2)}，
r2 で初めて max{A(l), ..., A(r2)} > min{B(l), ..., B(r2)}
となる l2 と r2 を2分探索で求めると，r2 - l2 + 1がそのときの答えになる．
あとは全ての l に対して足し合わせる．

計算量は O(NlogN)

ソースコード

#include <bits/stdc++.h>
using namespace std;

using ll = long long;

struct rmq1 { // min
    using type = int;
    static bool comp(type const& l, type const& r) {
        return l < r;
    }
};

struct rmq2 { // max
    using type = int;
    static bool comp(type const& l, type const& r) {
        return l > r;
    }
};

template<typename Policy>
class rmq_sparse_table {
    using T = typename Policy::type;

public:
    rmq_sparse_table(std::vector<T> const& v)
        : a(v), n(v.size()), log_n(v.size()+1)
    {
        for(int i=2; i<n+1; ++i) {
            log_n[i] = log_n[i >> 1] + 1;
        }

        table = std::vector<std::vector<int>>(n, std::vector<T>(log_n[n] + 1));
        for(int i=0; i<n; ++i) {
            table[i][0] = i;
        }
        for(int k=1; (1<<k) <= n; ++k) {
            for(int i=0; i + (1 << k) <= n; ++i) {
                int first = table[i][k-1],
                    second = table[i + (1 << (k-1))][k-1];
                if(Policy::comp(a[first], a[second])) {
                    table[i][k] = first;
                } else {
                    table[i][k] = second;
                }
            }
        }
    }

    // [l..r]
    int query(int l, int r) {
        int k = log_n[r - l + 1];
        if(Policy::comp(a[table[l][k]], a[table[r - (1 << k) + 1][k]])) {
            return table[l][k];
        } else {
            return table[r - (1 << k) + 1][k];
        }
    }

private:
    const int n;
    std::vector<int> log_n;
    std::vector<T> a;
    std::vector<std::vector<int>> table;
};


int main() {
    int N;
    cin >> N;
    vector<int> A(N), B(N);
    for(int i=0; i<N; ++i) {
        scanf("%d", &A[i]);
    }
    for(int i=0; i<N; ++i) {
        scanf("%d", &B[i]);
    }
    rmq_sparse_table<rmq1> st1(B);
    rmq_sparse_table<rmq2> st2(A);
    ll res = 0;
    for(int i=0; i<N; ++i) {
        if(A[i] > B[i] || A[st2.query(i, N-1)] < B[st1.query(i, N-1)]) {
            continue;
        }
        int l = i-1, r = N-1;
        while(r - l > 1) {
            int m = (r + l) / 2;
            if(A[st2.query(i, m)] < B[st1.query(i, m)]) {
                l = m;
            } else {
                r = m;
            }
        }
        const int l2 = r;
        if(A[st2.query(i, l2)] != B[st1.query(i, l2)]) {
            continue;
        }

        l = i, r = N;
        while(r - l > 1) {
            int m = (l + r) / 2;
            if(A[st2.query(i, m)] <= B[st1.query(i, m)]) {
                l = m;
            } else {
                r = m;
            }
        }
        const int r2 = l;
        res += (r2 - l2 + 1);
    }

    cout << res << endl;
}