永続セグ木

segtree.java

import java.util.function.Function;
import java.util.function.IntFunction;
import java.util.function.BinaryOperator;

/* セグ木 */
interface SegTree<T, S extends BinaryOperator<T>>
{
	/* 要素数 */
	int size();
	/* 
	 * 添字[l, r)の要素を半群の演算子で畳み込んだ結果を返す
	 * 添字の範囲が[0, size())に収まっていなければ例外を投げる
	 */
	T query(int l, int r) throws IndexOutOfBoundsException;
	/* 
	 * 添字iの要素x_iをf(x_i)で更新する
	 * iが[0, size())に収まっていなければ例外を投げる
	 */
	SegTree<T, S> update(int i, Function<T, T> f)
		throws IndexOutOfBoundsException;

	/* 要素全てを半群の演算子で畳み込んだもの */
	T data();
}

/* セグ木の葉を作るやつ */
/* いちいち添字の検査をするのは無駄なので，ここではしない */
final class SegTreeLeaf<T, S extends BinaryOperator<T>> implements SegTree<T, S>
{
	/* この葉が保持する要素 */
	private final T x;

	public SegTreeLeaf(T y) { x = y; }

	@Override public T data() { return x; }
	@Override public int size() { return 1; }

	@Override public T query(int l, int r) { return x; }
	@Override public SegTreeLeaf update(int i, Function<T, T> f)
	{ return new SegTreeLeaf(f.apply(x)); }
}

/* セグ木のノードを作るやつ */
/* ここでも添字の検査はしない */
final class SegTreeNode<T, S extends BinaryOperator<T>> implements SegTree<T, S>
{
	/* このノード以下の木にいくつ葉が含まれるか(=セグ木の要素数) */
	private final int n;
	/* このノードが保持する要素
	   (=セグ木の要素全てを半群の演算子で畳み込んだもの) */
	private final T x;
	/* 半群の演算子 */
	private final S op;
	/* このノードの子 */
	private final SegTree<T, S> t1, t2;

	public SegTreeNode(S f, SegTree<T, S> s, SegTree<T, S> t)
	{
		op = f;
		t1 = s;
		t2 = t;
		n = s.size() + t.size();
		x = f.apply(s.data(), t.data());
	}

	@Override public T data() { return x; }
	@Override public int size() { return n; }

	@Override public T query(int l, int r)
	{
		if (l == 0 && r == n) {
			return x;
		} else if (r <= t1.size()) {
			return t1.query(l, r);
		} else if (t1.size() <= l) {
			return t2.query(l - t1.size(), r - t1.size());
		} else {
			return op.apply(t1.query(l, t1.size()), t2.query(0, r - t1.size()));
		}
	}

	@Override public SegTreeNode update(int i, Function<T, T> f)
	{
		if (i < t1.size()) {
			return new SegTreeNode(op, t1.update(i, f), t2);
		} else {
			return new SegTreeNode(op, t1, t2.update(i - t1.size(), f));
		}
	}
}

/* 添字についての関数からセグ木を作るやつ（ここで添字の検査も入れる） */
final class SegTreeBuilder<T, S extends BinaryOperator<T>> implements SegTree<T, S>
{
	private final SegTree<T, S> t;

	private SegTreeBuilder(SegTree<T, S> t) { this.t = t; }

	private SegTree<T, S> segTreeBuilderAux(S op, int offset, int n, IntFunction<T> f)
	{
		if (n == 1) {
			return new SegTreeLeaf(f.apply(offset));
		} else {
			return new SegTreeNode(op, segTreeBuilderAux(op, offset, n / 2, f), segTreeBuilderAux(op, offset + n / 2, (n + 1) / 2, f));
		}
	}

	/* 半群の演算子，要素数，添字についての関数からセグ木を作るコンストラクタ */
	public SegTreeBuilder(S op, int n, IntFunction<T> f)
		throws IllegalArgumentException
	{
		if (n <= 0) {
			throw new IllegalArgumentException();
		} else {
			t = segTreeBuilderAux(op, 0, n, f);
		}
	}

	@Override public T data() { return t.data(); }
	@Override public int size() { return t.size(); }

	@Override public T query(int l, int r) throws IndexOutOfBoundsException
	{
		if (l < 0 || t.size() <= r) {
			throw new IndexOutOfBoundsException();
		} else {
			return t.query(l, r);
		}
	}

	@Override public SegTreeBuilder update(int i, Function<T, T> f)
		throws IndexOutOfBoundsException
	{
		if (i < 0 || t.size() <= i) {
			throw new IndexOutOfBoundsException();
		} else {
			return new SegTreeBuilder(t.update(i, f));
		}
	}
}