AOJ-2268, UTPC2011-J : Randomized Self-Balancing Binary Search Tree

問題概要

ノード数N(<30000)のTreapの高さがhになる確率を各hについて求める問題。

解法とか

コンテスト直後に解説放送を見ていたので打ち切り+FFTで解くことは知っていた。とはいえDPの式立てるのにも大分苦労したけど(高さh以下、というのを状態にするのに中々気づけなかった)。
D言語の標準ライブラリにFFTが入っていることを知って解こうと思った。実際に書いてみるとライブラリバグってて動かなかったし、ちょっと修正して使ってみたらspaghetti sourceの実装に比べて誤差が大きいらしくて大きいケースで盛大に誤差死していた。速度については、まあ当然C++での実装に比べると遅かったけど誤差が積もってnanが出ているせいかもしれないので何とも言えない。そんなに速くはないだろうと予想する。

acceptしたかったコード

import std.stdio;
import std.numeric;
import std.range;
import std.algorithm;
import std.math;
import core.bitop;
import std.complex;

int N, NN;

void init(){
	scanf("%d", &N);
}

double[] calc(double[] prev){
	auto xs = fft(prev), zs = new Complex!(double) [](NN);
	foreach(i, x; xs){
		zs[i] = x * x;
	}
	zs = inverseFft(zs);
	auto ret = new double[](NN);
	fill(ret, 0.0);
	ret[0] = 1.0;
	foreach(i ; 1..N+1){
		ret[i] = zs[i-1].re / i;
	}
	return ret;
}

void solve(){
	NN = 4<<bsr(N);
	auto prev = new double[](NN);
	fill(prev, 0.0);
	prev[0] = prev[1] = 1.0;
	bool small = false;
	writeln(prev[N]);
	for(int i=0; i<N-1; ++i){
		if(small){
			writeln("0");
		}
		else{
			auto cur = calc(prev);
			writefln("%.12f", cur[N] - prev[N]);
			if(i > 50 && cur[N] - prev[N] < 1e-6){
				small = true;
			}
			prev = cur;
		}
	}
}

void main(){
	init;
	solve;
}