题目链接: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);
}
}