読者です 読者をやめる 読者になる 読者になる

POJ 3321 Apple Tree (日記を兼ねる)

今日はジョギングしたら眠くなったのでコンテスト前に仮眠しようと布団に入って、起きたらAtcoerが終わっていた。ショックだった。
日記はここで終わり。

http://poj.org/problem?id=3321

根付き木になるようにDFSすると部分木が連番になるので、その部分の和をFenwick Treeを使って持てる、という問題。
わりと面白かったけどTLEが無慈悲すぎる。
6TLEで断念しました。

//8WA 6TLE

#include <cstdio>
#include <iostream>
#include <vector>
#include <algorithm>
#include <cstring>
#include <bitset>
using namespace std;
#define REP(i,s,n) for(int i=(int)(s);i<(int)(n);++i)
#define rep(i,n) REP(i, 0, n)

template<class T, int MAX_N>
struct FenwickTree { //0-indexed
    int N;
    T data[MAX_N];
    void init(int n) {
        N = n;
        rep(i, N) data[i] = (i + 1) & (-i - 1);
    }
    T sum(int j) { //[0, j)
        T ret = 0;
        for(int x = j - 1; x >= 0; x = (x & (x + 1)) - 1)
          ret += data[x];
        return ret;
    }
    void add(int j, T w) { // 0 <= j < N
        for(int x = j; x < N; x |= x + 1)
          data[x] += w;
    }
};

vector<int> G[100000];
int treeID[100000];
int max_range[100000]; // treeID[i] のiに対しクエリに使うrangeを返す

int cnt_nodes;
void dfs(int cur_v, int p) {
    treeID[cur_v] = cnt_nodes++;
    rep(i, G[cur_v].size())
      if(G[cur_v][i] != p)
        dfs(G[cur_v][i], cur_v);
    max_range[cur_v] = cnt_nodes - 1;
}

bool apple[100000];
int main() {
    int N; scanf("%d", &N);
    rep(_, N - 1) {
        int u, v;
        scanf("%d%d", &u, &v);
        --u; --v;
        G[u].push_back(v);
        G[v].push_back(u);
    }
    cnt_nodes = 0;
    dfs(0, 0);
    int M; scanf("%d", &M);
    FenwickTree<int, 100000> fwt;
    fwt.init(N);
    rep(_, M) {
        char q; int v;
        scanf(" %c%d", &q, &v);
        --v;
        if(q == 'Q') {
            printf("%d\n", fwt.sum(max_range[v] + 1) - fwt.sum(treeID[v]));
        }
        else {
            fwt.add(treeID[v], apple[v] ? 1 : -1);
            apple[v] ^= 1;
        }
    }
    return 0;
}

なんでこの問題八千人もACしてるのか全くわかりません。
C++11ならemplace_back使うとかC++14ならfenwick treeの初期化にconstexpr使うとかもアリだと思うんですがPOJはc++03しか使えないしなんか異様に遅いしこれ以上は入出力いじらないと無理っぽいような。

Fenwick treeは毎回蟻本を写経していたので、http://hos.ac/slides/20140319_bit.pdfを参考にコピペして使えそうなのを作りました。
[追記 (2017/2/15)]
某所で隣接リストにvectorのかわりにpairの配列をソートしたものを使うと速いと聞いたので試したら通った。

#include <cstdio>
#include <iostream>
#include <vector>
#include <algorithm>
#include <cstring>
#include <bitset>
using namespace std;
#define REP(i,s,n) for(int i=(int)(s);i<(int)(n);++i)
#define rep(i,n) REP(i, 0, n)

template<class T, int MAX_N>
struct FenwickTree { //0-indexed
    int N;
    T data[MAX_N];
    void init(int n) {
        N = n;
        rep(i, N) data[i] = (i + 1) & (-i - 1);
    }
    T sum(int j) { //[0, j)
        T ret = 0;
        for(int x = j - 1; x >= 0; x = (x & (x + 1)) - 1)
          ret += data[x];
        return ret;
    }
    void add(int j, T w) { // 0 <= j < N
        for(int x = j; x < N; x |= x + 1)
          data[x] += w;
    }
};

int N;
pair<int, int> edges[200000];
int vToe[100000];
int treeID[100000];
int max_range[100000]; // treeID[i] のiに対しクエリに使うrangeを返す

int cnt_nodes;
void dfs(int cur_v, int p) {
    treeID[cur_v] = cnt_nodes++;
    if(vToe[cur_v] != -1) {
        for(int i = vToe[cur_v]; i < 2 * (N - 1); ++i) {
            if(edges[i].first != cur_v) break;
            if(edges[i].second != p)
                dfs(edges[i].second, cur_v); 
        }
    }
    max_range[cur_v] = cnt_nodes;
}

bool apple[100000];
FenwickTree<int, 100000> fwt;

int main() {
    scanf("%d", &N);
    rep(i, N - 1) {
        int u, v;
        scanf("%d%d", &u, &v);
        --u; --v;
        if(v > u) swap(u, v);
        edges[i] = make_pair(u, v);
        edges[i + N - 1] = make_pair(v, u);        
    }
    sort(edges, edges + 2 * (N - 1));
    memset(vToe, -1, sizeof(vToe));
    
    for(int i = 0; i < 2 * (N - 1); ++i)
        if(i == 0 || edges[i].first != edges[i - 1].first)
            vToe[edges[i].first] = i;
    
    cnt_nodes = 0;
    dfs(0, 0);
    int M; scanf("%d", &M);
    fwt.init(N);
    rep(_, M) {
        char q; int v;
        scanf(" %c%d", &q, &v);
        --v;
        if(q == 'Q') {
            printf("%d\n", fwt.sum(max_range[v]) - fwt.sum(treeID[v]));
        }
        else {
            fwt.add(treeID[v], apple[v] ? 1 : -1);
            apple[v] ^= 1;
        }
    }
    return 0;
}