Find Kth Largest XOR Coordinate Value LeetCode solution in Java

1
class Solution {
2

3
  private static final Random RAND = new Random(0);
4

5
  public int kthLargestValue(int[][] matrix, int k) {
6
    var xor = convertToXorArray(matrix);
7
    var targetIdx = xor.length - k;
8
    sortPartially(xor, targetIdx, 0, xor.length);
9
    return xor[targetIdx];
10
  }
11

12
  void sortPartially(int[] nums, int targetIdx, int origLo, int origHi) {
13
    if (origHi - origLo < 2) return;
14

15
    var pivotIdx = RAND.nextInt(origHi - origLo) + origLo;
16
    var pivot = nums[pivotIdx];
17
    swap(nums, origLo, pivotIdx);
18

19
    var lo = origLo;
20
    var mid = lo + 1;
21
    var hi = origHi;
22
    while (mid < hi) {
23
      if (pivot < nums[mid]) swap(nums, mid, --hi);
24
      else if (pivot > nums[mid]) swap(nums, mid++, lo++);
25
      else mid++;
26
    }
27

28
    if (targetIdx < lo) sortPartially(nums, targetIdx, origLo, lo);
29

30
    sortPartially(nums, targetIdx, mid, origHi);
31
  }
32

33
  void swap(int[] n, int p, int q) {
34
    var tmp = n[p];
35
    n[p] = n[q];
36
    n[q] = tmp;
37
  }
38

39
  int[] convertToXorArray(int[][] matrix) {
40
    var rows = matrix.length;
41
    var cols = matrix[0].length;
42
    var xor = new int[rows * cols];
43
    for (int r = 0; r < rows; r++) {
44
      for (int c = 0; c < cols; c++) {
45
        var xIdx = r * cols + c;
46
        xor[xIdx] = matrix[r][c];
47
        if (c > 0) xor[xIdx] ^= xor[xIdx - 1];
48
      }
49
    }
50

51
    for (int i = cols; i < xor.length; i++) xor[i] ^= xor[i - cols];
52

53
    return xor;
54
  }
55
}