yukicoder No.59 - 鉄道の旅

yukicoder Fenwick tree

問題

http://yukicoder.me/problems/no/59

解法

Fenwick Treeを使うだけ.

ソースコード

#include <iostream>
#include <vector>
using namespace std;
/*
 * Fenwick Treeの実装.
 * 区間[1, n]において, いわゆるfoldlの操作を高速化することができる.
 */
template<typename T>
class FenwickTree {
public:
    vector<T> elt_;
    int n_;

    FenwickTree() { }

    FenwickTree(int n) {
        elt_ = vector<T>(n + 1);
        n_ = n;
        /* 初期化処理 */
        for (int i = 0; i <= n_; i++) {
            elt_[i] = 0;
        } 
    }

    void add(int i, T x) {
        while (i <= elt_.size()) {
            elt_[i] += x;
            i += i & -i;
        }
    }

    T query(int i) {
        T s = 0;
        while (i > 0) {
            /* クエリの処理 */
            s += elt_[i];
            i -= i & -i;
        }
        return s;
    }
};

const int kMAX_N = 100010;
const int kMAX_W = 1000010;

int N, K;
int W[kMAX_N];

void Solve() {
    FenwickTree<int> tree = FenwickTree<int>(kMAX_W);

    for (int i = 0; i < N; i++) {
        if (W[i] > 0 && tree.query(kMAX_W) - tree.query(W[i] - 1) < K) {
            tree.add(W[i], 1);
        } else if (W[i] < 0 && tree.query(abs(W[i])) - tree.query(abs(W[i]) - 1) > 0) {
            tree.add(abs(W[i]), -1);
        }
    }
    cout << tree.query(kMAX_W) << endl;
}

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);

    cin >> N >> K;
    for (int i = 0; i < N; i++) {
        cin >> W[i];
    }

    Solve();

    return 0;
}

感想

重さWの荷物を何個積んでいるかの配列を作っておけばもう少しシンプルになります. あとFenwick Treeって名前かっこいい.