题目:
You are given an integer array nums and you have to return a new counts array. The counts array has the property where counts[i] is the number of smaller elements to the right of nums[i]. Example: Given nums = [5, 2, 6, 1] To the right of 5 there are 2 smaller elements (2 and 1). To the right of 2 there is only 1 smaller element (1). To the right of 6 there is 1 smaller element (1). To the right of 1 there is 0 smaller element. Return the array [2, 1, 1, 0].
分析:很显然不能用时间复杂度为O(N2)的暴力穷举法,下面先给出一个运行时间为102ms的AC答案:
public class Solution { private void add(int[] bit, int i, int val) { for (; i < bit.length; i += 1) bit[i] += val; } private int query(int[] bit, int i) { return bit[i]; } public List<Integer> countSmaller(int[] nums) { int[] tmp = nums.clone(); Arrays.sort(tmp); for (int i = 0; i < nums.length; i++) { nums[i] = Arrays.binarySearch(tmp, nums[i]); } int[] bit = new int[nums.length]; Integer[] ans = new Integer[nums.length]; for (int i = nums.length - 1; i >= 0; i--) { ans[i] = query(bit, nums[i]); add(bit, nums[i]+1, 1); } return Arrays.asList(ans); } }
上面的答案虽然AC了,但若原始海量数据本身是降序排列时,再运行时间也为O(N2),下面给出运用树状数组(Binary Indexed Tree (Fenwick tree))的解法,使时间复杂度将为O(NlgN),树状数组的相关知识可参考博客树状数组.
private int lowbit(int x) { return x&(-x); } private void add(int[] bit, int i, int val) { for (; i < bit.length; i += lowbit(i)) bit[i] += val; } private int query(int[] bit, int i) { int ans = 0; for (; i > 0; i -= lowbit(i)) ans += bit[i]; return ans; } public List<Integer> countSmaller(int[] nums) { int[] tmp = nums.clone(); Arrays.sort(tmp); for (int i = 0; i < nums.length; i++) nums[i] = Arrays.binarySearch(tmp, nums[i]); int[] bit = new int[nums.length]; Integer[] ans = new Integer[nums.length]; for (int i = nums.length - 1; i >= 0; i--) { ans[i] = query(bit, nums[i]); add(bit, nums[i]+1, 1); } return Arrays.asList(ans); }
最后给出线段树(segment tree)的解法,时间复杂度也为O(NlgN),相似题目请参考Range Sum Query - Mutable,代码如下:
class SegmentTreeNode { int start, end; int num; SegmentTreeNode ltree, rtree; public SegmentTreeNode(int s, int e) { start = s; end = e; } } public class Solution { SegmentTreeNode root = null; public SegmentTreeNode buildTree(int[] nums, int left, int right) { SegmentTreeNode root = new SegmentTreeNode(left, right); if (left != right) { int mid = left + (right - left)/2; root.ltree = buildTree(nums, left, mid); root.rtree = buildTree(nums, mid+1, right); } return root; } private void update(SegmentTreeNode root, int i, int val) { if (root.start == root.end) { root.num += 1; } else { int mid = root.start + (root.end - root.start)/2; if (i <= mid) { update(root.ltree, i, val); } else { update(root.rtree, i, val); } root.num = root.ltree.num + root.rtree.num; } } private int query(SegmentTreeNode root, int i, int j) { if (root.start == i && root.end == j) { return root.num; } else { int mid = root.start + (root.end - root.start)/2; if (j <= mid) { return query(root.ltree, i, j); } else if (i > mid) { return query(root.rtree, i, j); } else { return query(root.ltree, i, root.ltree.end) + query(root.rtree, root.rtree.start, j); } } } public List<Integer> countSmaller(int[] nums) { int[] tmp = nums.clone(); Arrays.sort(tmp); for (int i = 0; i < nums.length; i++) { nums[i] = Arrays.binarySearch(tmp, nums[i]) + 1; } int[] bit = new int[nums.length + 1]; root = buildTree(bit, 0, bit.length-1); Integer[] ans = new Integer[nums.length]; for (int i = nums.length - 1; i >= 0; i--) { ans[i] = query(root, 0, nums[i] - 1); update(root, nums[i], 1); } return Arrays.asList(ans); } }