SWAG(Sliding Window Aggregation)の実装についての提案

注意

本記事は筆者が昔若気のいたりで書いた記事です。有用性がないわけでもないので、読むなとまでは言いませんが、新たに書き直した記事も合わせてご覧ください。
motsu-xe.hatenablog.com

注意終わり

〜〜〜〜〜〜〜〜〜〜〜〜〜〜〜〜〜〜〜

おはようございます。この業界では何時でもおはようございますが挨拶です。

今回は少し前にTwitterでこの話題持ちきりになっていた、SWAGについてのお話をしたいと思います。とはいえ私はデータ構造に造詣が深いわけではないので、あまり深くに突っ込むことはできません。一応紹介はしますが、飽くまでも実用的な実装の話をしたいだけなので、詳しくは他の方の記事をご覧ください。

SWAGとは

スライド最小値を知っていますか？ある制約を満たした複数の区間に対して、その区間最小値を求める高速に求めるデータ構造です。詳しくはググって欲しいのですが、簡単に言うと以下に示す通りです。

全順序集合の $N$ 要素の配列 $v$ が与えられる。 $M$ 個の区間 $[s_i,t_i)\ (0\leq i \lt m,\ 0\leq s_i \lt t_i \lt n)$ での配列要素の最小値を求める。この操作全体を $O(N+M)$ で行う。ただし区間列は $1\leq\forall i\lt n,\,\ s_{i-1}\leq s_i,\ t_{i-1}\leq t_i$ を満たしている必要がある。

尺取りのように区間の左端と右端をスライドさせながら、最小値を求めるのでスライド最小値と呼ばれます。
SWAGと言うのは可能な操作だけを見れば、スライド最小値の完全上位互換のデータ構造です。全順序集合における $min$ 演算について区間での演算結果を返すのがスライド最小値なのに対し、モノイドにおける演算について区間での演算結果を返すのがSWAGです。一般に全順序集合に $min$ で演算を入れたものはモノイドになる*1ので、一般化になっていることがわかります。一応要件を整理しておきましょう。

モノイドの $N$ 要素の配列 $v$ が与えられる。 $M$ 個の区間 $[s_i,t_i)\ (0\leq i \lt m,\ 0\leq s_i \lt t_i \lt n)$ での配列要素の演算結果を求める。この操作全体を $O(N+M)$ *2で行う。ただし区間列は $1\leq\forall i\lt n,\,\ s_{i-1}\leq s_i,\ t_{i-1}\leq t_i$ を満たしている必要がある。

では次にどのような方法でこの操作を実現するかを紹介します。

仕組み

基本的にtwo-stack-queueを用いることになります。frontスタックとbackスタックという2つのstackを用意します。これらのスタックはモノイドのペアを持ちます。

区間の右端を進めるとき

進めた分の配列の要素を、frontスタックに詰め込みます。このとき、ペアの第2要素にはfrontスタックに入っている要素の総演算結果を持っておきます。これは右端を進める前のfrontスタックの先頭要素の第2要素に、新たに詰める要素を演算することで計算可能です。
自然数と積に関するモノイドに関して例示すると

末尾↔︎先頭
[(2,2),(3,6)]
↓(5を追加)
[(2,2),(3,6),(5,30)]

のようになります。

区間の左端を進めるとき

backスタックから1つ要素をpopします。
同様に例示すると

先頭↔︎末尾
[(3,6),(2,2)]
↓(3を除外)
[(2,2)]

ただし、backスタックが空のときは、frontスタックの要素を1つずつbackスタックに移し、その後backスタックから1つ要素をpopします。移す際はfrontスタックと同様に、ペアの第2要素にはbackスタックに入っている要素の総演算結果を持っておきます。
例示すると

back	front
先頭↔︎末尾	末尾↔︎先頭
[]	[(2,2),(3,6)]
↓(3を移す)
[(3,3)]	[(2,2)]
↓(2を移す)
[(2,6),(3,3)]	[]
↓(2を除外)
[(3,3)]	[]

となる。

演算結果を計算するとき

区間を定め終わったときに、演算結果を計算するにはfrontとbackの先頭要素の第2要素同士を演算すれば良いです。今演算したいものはfrontないしbackに入っており、それぞれの先頭要素の第2要素は、それぞれのスタックに今入っている要素の総演算結果が入っているため、その2つを演算することで、全体の総演算結果が求められます。もしもスタックが空だった場合は、単位元とすれば良いです。

計算量

配列の各要素につき、frontにpush、frontからbackに移動、backからpopしか行われないため、 $O(N+M)$ を達成しています。

以上がSWAGの説明となります。これでわかりますか？わからなかったと思うので、参考資料の方を読みましょう。

提案

ここからが本題(すぐ終わる)なのですが、上の方針を最初に見たときに、結構無駄があるなぁと言うのと、スライド最小値のスライドしている感が分かり辛いかなぁと言うのとを感じたので、自分なりにアレンジして実装してみました。
まず最初に、使用用途を考えたとき、基本的に配列上をスライドしながら最小値を求める、と言う形になることが多いと思いました。なのでqueueのような操作よりも、最初に述べたような区間クエリの操作として見れた方が、わかりやすいと感じました。
そうなった際に、配列を持っておくと、かなり無駄が多い事に気付きます。まずfrontスタックですが、明示的に持つ必要が無くなります。実際使っているのは、frontの先頭要素の第2要素だけであることがわかります。配列を別に持っているため、indexだけ覚えておけば、後々backスタックに突っ込むときにも困りません。よって実際に持つ必要があるのは、frontスタックの総演算結果のみとなります。
またbackスタックですが、使うのは演算結果のみであるので、それぞれの要素を覚えておく必要は全くありません。よって詰め込むのは演算結果のみとすることができます。
以上を踏まえて実装したものが、以下となります。

template <class Monoid,Monoid e> struct SWAG{
    std::function<Monoid(Monoid,Monoid)> op;
    std::vector<Monoid> v;
    Monoid front;
    std::stack<Monoid> back;
    size_t l,r;
    explicit SWAG(std::function<Monoid(Monoid,Monoid)> op)
            :op(std::move(op)),v(),front(e),back(),l(0),r(0){};
    SWAG(std::function<Monoid(Monoid,Monoid)> op,std::vector<Monoid>&v_)
            :op(std::move(op)),v(v_.size()),front(e),back(),l(0),r(0){
        std::copy(v_.begin(),v_.end(),v.begin());
    };
    Monoid fold(size_t i,size_t j){
        while(r<j) front=op(front,v[r++]);
        while(l<i){
            if(back.empty()){
                Monoid tmp{e};
                for(size_t u=r;u-->l;) back.emplace(tmp=op(v[u],tmp));
                front=e;
            }
            back.pop();
            ++l;
        }
        if(back.empty()) return front;
        return op(back.top(),front);
    }
};

ごちゃごちゃしていますし、可読性もだいぶ低い実装をしていますが、fold関数のところだけみてくれれば大体わかると思います。話したいことは以上となります。

使用例

www.hackerrank.com

提出コード

#include <iostream>
#include <functional>
#include <vector>
#include <stack>

//SWAGは省略

template<class T>
T gcd(T n,T m){
    while(m){
        n%=m;
        std::swap(n,m);
    }
    return n;
}

template<class T>
T lcm(T n,T m){
    T c=n*m/gcd(n,m);
    if(c>1000000001) c=1000000001;
    return c;
}

int main(){
    using ll = long long;
    int n;
    ll x;
    std::cin>>n>>x;
    std::vector<ll> a(n);
    for(auto&i:a) std::cin>>i;
    std::function<ll(ll,ll)> op=[](ll i,ll j){return lcm(i,j);};
    SWAG<ll,1> sw(op,a);
    int r=0,ans=0,ansl=0;
    for(int l=0;l<n;++l){
        while(r<n&&sw.fold(l,r+1)<x) ++r;
        if(ans<r-l){
            ans=r-l;
            ansl=l;
        }
    }
    std::cout<<ansl+1<<" "<<ansl+ans<<std::endl;
}