Hackerrank - Similar Pair 与线段树解法

题目链接：https://www.hackerrank.com/challenges/similarpair

一开始想到的简单解法如下（最后两个testCase超时）：

package similarpair;

public class Solution {
	static java.util.Scanner in = new java.util.Scanner(System.in);
	static java.io.PrintStream out = System.out;
	static int count;
	public static final int NULLID = 0;

	public static void main(String[] args) {
		int n = in.nextInt();
		int t = in.nextInt();
		/**
		 * parent[i] 表示序号为i的节点的父节点序号
		 */
		int[] parents = new int[n + 1];
		count = 0;
		for (int i = 1; i < n; i++) {
			int parentId = in.nextInt();
			int childId = in.nextInt();
			parents[childId] = parentId;
		}
		for (int i = 1; i <= n; i++) {
			int id = parents[i];
			while (id != NULLID) {
				if (Math.abs(id - i) <= t) {
					count++;
				}
				id = parents[id];
			}
		}
		out.println(count);
	}

}

于是又去网上搜索别人的解法。

搜到的几篇都是只给代码，不给思路，潦草的看过去，完全没能理解。之后耐着性子仔细地读他的代码才理解。

代码被我放到Eclipse里用Refactor功能改了一大片函数参数个数与名字，将SegmentTree与Solution类分离开来，期间发现了一些奇妙的技巧。加深了对SegmentTree的印象。

重构与注解后的代码如下：

/*
original java source code from http://blog.csdn.net/tspatial_thunder/article/details/42832687
 */
package similarpair_st;

import java.util.*;

public class Solution {

	@SuppressWarnings("unchecked")
	public static LinkedList<Integer>[] children = new LinkedList[100002];
	static int n, t, rootId;
	static SegmentTree segmentTree;
	static long pairs;

	public static void main(String[] args) {

		@SuppressWarnings("resource")
		Scanner scan = new Scanner(System.in);

		n = scan.nextInt();
		t = scan.nextInt();
		segmentTree = new SegmentTree(n);

		for (int i = 1; i <= n; i++)
			children[i] = new LinkedList<Integer>();
		// 记录每个节点的入度，根据入度是否为0来找出根节点
		int[] inDegrees = new int[n + 1];

		for (int i = 1; i < n; i++) {
			int parentId = scan.nextInt();
			int childId = scan.nextInt();

			children[parentId].add(childId);
			inDegrees[childId]++;
		}

		for (int i = 1; i <= n; i++) {
			if (inDegrees[i] == 0) {
				rootId = i;
				break;
			}
		}

		pairs = 0L;
		depthSearch(rootId);

		System.out.println(pairs);

	}
	/*
	 * 巧妙地结合递归搜索，在线段树上更新“当前节点是之后所有后续深入下去的节点的父节点”这条信息。
	 * query操作是标准的范围查询。update则是从根开始，更新“有多少父节点”这条信息。
	 * 至于递归深入前后的两次互逆的update操作，则是为了消去“当前节点是后续节点的父节点”信息。
	 * 因为递归的回溯操作，意味着过去记录的“是父节点”信息，对于另一条分支，其实并不是父节点，所以要消去。
	 */
	public static void depthSearch(int nodeId) {

		int min = (nodeId - t < 1) ? 1 : nodeId - t;
		int max = (nodeId + t > n) ? n : nodeId + t;

		pairs += segmentTree.query(min, max);

		segmentTree.update(nodeId, 1);

		for (int childId : children[nodeId]) {
			depthSearch(childId);
		}

		segmentTree.update(nodeId, -1);
	}
}

/**
 * 此线段树的所分割的线段，其左端点为1，其右端点为segmentLength。
 */
class SegmentTree {
	int segmentLength;
	/**
	 * 此线段树的节点只有一个数据域，类型为long，其含义是在该节点表示的线段范围里，有多少是满足题意的上级节点。
	 * 节点缺失的线段起点与线段终点隐式地存在于query和update方法的segmentBegin、segmentEnd参数中。
	 * 即每次query、update方法调用时给出的nodeIndex所指示的节点，其线段的左右端点分别是segmentBegin、segmentEnd。
	 * 所以不需要像这样子：
	 * class Node { long data; int left; int right; }
	 */
	long[] tree;

	public SegmentTree(int segmentLength) {
		this.segmentLength = segmentLength;
		tree = new long[4 * segmentLength + 1];
	}

	public static final int ROOT_INDEX = 1;

	public static int parentIndex(int childIndex) {
		return childIndex >> 1;
	}

	public static int leftChildIndex(int parentIndex) {
		return parentIndex << 1;
	}

	public static int rightChildIndex(int parentIndex) {
		return (parentIndex << 1) | 1;
	}

	public long query(int begin, int end) {
		return query(ROOT_INDEX, 1, segmentLength, begin, end);
	}

	public long query(int nodeIndex, int segmentBegin, int segmentEnd, int begin, int end) {

		if (end < segmentBegin || begin > segmentEnd)
			return 0;

		else if (end == segmentEnd && begin == segmentBegin)
			return tree[nodeIndex];

		else {
			int mid = (segmentBegin + segmentEnd) >> 1;
			int lmax = (mid < end) ? mid : end;
			int rmin = (begin > mid) ? begin : mid + 1;
			return query(leftChildIndex(nodeIndex), segmentBegin, mid, begin, lmax)
					+ query(rightChildIndex(nodeIndex), mid + 1, segmentEnd, rmin, end);
		}

	}

	public void update(int value, long opt) {
		update(ROOT_INDEX, 1, segmentLength, value, opt);
	}

	public void update(int nodeIndex, int segmentBegin, int segmentEnd, int val, long opt) {
		if (val < segmentBegin || val > segmentEnd || segmentBegin > segmentEnd)
			return;

		tree[nodeIndex] += opt;

		int m = (segmentBegin + segmentEnd) >> 1;

		if (segmentBegin == segmentEnd)
			return;
		else if (val <= m)
			update(leftChildIndex(nodeIndex), segmentBegin, m, val, opt);
		else
			update(rightChildIndex(nodeIndex), m + 1, segmentEnd, val, opt);
	}
}