POJ-3321: Apple Tree

keyword

BIT　木　定数倍最適化　C++ Java

問題概要

ノード数N(<10^5)の根付き木が与えられる。M(<10^5)のクエリを処理する問題。クエリは、あるノードの値が増減したという情報と、あるノードを根とする部分木の値の和を出力する2種類である。

解法

根から行きがけ順に番号を降ると、子孫のノードには連続する範囲の整数が割り振られる。したがって、後はBITで区間和を取り出してやれば良い。計算量O(N + M*log N)。

付記

TLEがアホみたいに厳しい。注意点として、dfsで書くと深さが10^5のオーダーで問題になるけど、それを回避しても(stack使って展開したり)なお厳しい。まだ通って無いけど方針はさすがにこれで間違えてないだろうと思う(log が落ちることはないだろう)。

さらに付記

定数倍最適化が厳しいのでJavaで解いた。変に聞こえるようだけど、Javaは実行時間に3倍以上のボーナスがついているのでそれ狙い。もちろんIOの多い問題では適当に高速化してやる必要がある。

import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.Scanner;

public class Main{

	Scanner in;
	PrintWriter out;

	int[] table, BIT, tos;
	boolean[] noApple;
	ArrayList<Integer>[] G;
	int N, counter;

	void run(){
		out = new PrintWriter(System.out);
		in = new Scanner(System.in);
		N = readInt();
		G = new ArrayList[N+1];
		table = new int[N+1];
		BIT = new int[N+1];
		tos = new int[N+1];
		noApple = new boolean[N+1];
		for(int i=0; i<G.length; i++){
			G[i] = new ArrayList<Integer>();
		}
		for(int i=0, x, y; i<N-1; i++){
			x = readInt(); y = readInt();
			G[x].add(y); G[y].add(x);
		}
		counter = 1;
		dfs(1);
		int M = readInt();
		for(int i=0, x; i<M; i++){
			String op = in.next();
			x = table[readInt()];
			if(op.equals("Q")){
				System.out.println(tos[x]-x-sum(x,tos[x]));
			}
			else{
				noApple[x] = !noApple[x];
				add(x, noApple[x]?1:-1);
			}
		}
		out.flush();
	}

	int sum(int n){
		int _R = 0;
		for(;n>0;n-=n&-n){
			_R += BIT[n];
		}
		return _R;
	}
	
	int sum(int from, int to){
		return sum(to-1) - sum(from-1);
	}

	void add(int n, int x){
		for(;n<=N; n+=n&-n){
			BIT[n] += x;
		}
	}

	void dfs(int v){
		int gv = G[v].size();
		table[v] = counter++;
		for(int i=0; i<gv; i++){
			int u = G[v].get(i);
			if(table[u]==0){
				dfs(u);
			}
		}
		tos[table[v]] = counter;
	}

	int readInt(){
		return Integer.parseInt(in.next()); 
	}

	public static void main(String[] args){
		new Main().run();
	}
}

TLE解。

#include <cstdio>
#include <vector>
using namespace std;

const int MAX_N = 100009;
int table[MAX_N+1];
int BIT[MAX_N+1];
int tos[MAX_N+1];
int idx[MAX_N+1];
bool noApple[MAX_N+1];
vector<int> G[MAX_N+1];
int N, counter;

inline int readint(){
    int ret = 0, x;
    while(x=getchar(), !('0'<=x&&x<='9'));
    ret = (x&15);
    while(x=getchar(), '0'<=x&&x<='9') ret = (ret<<3) + (ret<<1) + (x&15);
    return ret;
}

inline int sum(int n){int _R=0;for(;n;n-=n&-n)_R+=BIT[n];return _R;}
inline int sum(int from, int to){return sum(to-1)-sum(from-1);}
inline void add(int n, int x){for(;n<=N;n+=n&-n)BIT[n]+=x;}

inline void dfs(int v){
	int gv = G[v].size();
	table[v] = counter++;
	for(int i=0; i<gv; i++){
		if(!table[G[v][i]]) dfs(G[v][i]);
	}
	tos[table[v]] = counter;
}

int main(){
	N = readint();
	for(int i=0, x, y; i<N-1; i++){
		x = readint(), y = readint();
		G[x].push_back(y); G[y].push_back(x);
	}
	counter = 1;
	dfs(1);
	int M = readint();
	for(int i=0, x; i<M; i++){
		char op;
		scanf(" %c ",&op);
		x = table[readint()];
		if(op=='Q'){
			printf("%d\n", tos[x]-x-sum(x,tos[x]));
		}
		else{
			noApple[x] = !noApple[x];
			add(x,noApple[x]?1:-1);
		}
	}
	return 0;
}