长度为N的数组arr,一定可以组成N^2个数值对。
例如arr = [3,1,2],
数值对有(3,3) (3,1) (3,2) (1,3) (1,1) (1,2) (2,3) (2,1) (2,2),
也就是任意两个数都有数值对,而且自己和自己也算数值对。
数值对怎么排序?规定,第一维数据从小到大,第一维数据一样的,第二维数组也从小到大。所以上面的数值对排序的结果为:
(1,1)(1,2)(1,3)(2,1)(2,2)(2,3)(3,1)(3,2)(3,3)
给定一个数组arr,和整数k,返回第k小的数值对。
public class Solution {
public static void swap(int[] arr, int i, int j) {
int tmp = arr[i];
arr[i] = arr[j];
arr[j] = tmp;
}
private static int[] partition(int[] nums, int L, int R) {
int pivot = nums[(int) (Math.random() * (R - L + 1)) + L];
int less = L - 1;
int more = R + 1;
int index = L;
while (index < more) {
if (nums[index] < pivot) {
swap(nums, ++less, index++);
} else if (nums[index] > pivot) {
swap(nums, index, --more);
} else {
index++;
}
}
return new int[]{less + 1, more - 1};
}
private static int getKth(int[] nums, int K) {
int L = 0, R = nums.length - 1;
while (L < R) {
int[] partition = partition(nums, L, R);
if (K >= partition[0] && K <= partition[1]) {
return nums[K];
} else if (K < partition[0]) {
R = partition[0] - 1;
} else {
L = partition[1] + 1;
}
}
return nums[L];
}
private static int[] solve(int[] nums, int K) {
int n = nums.length;
int first = getKth(nums, (K - 1) / n);
int lessNum = 0;
int equalNum = 0;
for (int num : nums) {
if (num < first) {
lessNum++;
} else if (num == first) {
equalNum++;
}
}
int rest = K - lessNum * n;
return new int[]{first, getKth(nums, (rest - 1) / equalNum)};
}
}