No.59 鉄道の旅

No.59 鉄道の旅 - yukicoder

• 典型的な RMQ っぽい問題。
• 蟻本で BIT の存在はしっていたけど使うのは初めてだった。
• この問題では sum は a[i]+...+a[maxW] なんだけど BIT では逆で a[1] + ... + a[i] 。それゆえに a[i] = sum[i]-sum[i+1] となる。これをうっかり逆にしてしまい WA した。

class RailroadTrip {
public:
    void solve(void) {
            int N,K;
            cin>>N>>K;
            // 単純にシミュレーションすればよい
            // 駅数が最大 10^5 なので積み荷の載せ・降ろしの判定を高速でできればよい。

            // 次のクエリを log(N) くらいでできればよい
            //  * a[w] += 1
            //  * a[w]+a[w+1]+...+a[maxW] を計算
            // BIT を使う

            int maxW = (1<<(int)ceil(log2(1E+6)));

            // [1,maxW]
            vector<ll> bits(maxW+1,0);
            // a[i] に x を追加する
            auto add = [&](int i, int x) {
                assert(i>0);
                while (i <= maxW)
                {
                    bits[i] += x;
                    i += (i & -i);
                }
            };
            // i 以上の a[i] の和
            auto sum = [&](int i) {
                ll s = 0;
                --i; // 求めたいのは (a[1]+...+a[maxW]) - (a[1]+...+a[i-1]) なので
                while (i > 0)
                {
                    s += bits[i];
                    i -= (i & -i);
                }
                return bits[maxW] - s;
            };
            // O(N*log(N))
            REP(i,N)
            {
                ll w, aw;
                cin>>w;
                aw = abs(w);
                if (w > 0 && sum(w) < K)
                    add(w,1);
                else if (w < 0 && sum(aw)-sum(aw+1) > 0)
                    add(aw,-1);
            }
            cout<<sum(1)<<endl;
    }
};