Kth Elements in some Sorted Arrays

简化版:https://www.geeksforgeeks.org/k-th-element-two-sorted-arrays/

用三元组记录坐标和value,很巧妙,不然要用class。

    public int mergeSort(int[][] nums, int k) {
        int rowNum = nums.length;
        PriorityQueue<int[]> pq = new PriorityQueue<>(rowNum, new Comparator<int[]>() {
            public int compare(int[] i1, int[] i2) {
                return i1[2] - i2[2];
            }
        });

        // 把每一行的第一个元素,加入到优先队列中,count记录加入了多少个元素
        int count = 0;
        //  sum用来数二维数组中一共有多少个元素
        int sum = 0;
        for (int i = 0; i < nums.length; i++) {
            if (nums[i].length > 0) {
                int[] arr = new int[]{i, 0, nums[i][0]};
                pq.offer(arr);
                count++;
                sum += nums[i].length;
            }
        }

        //  如果元素个数小于k,说明第k小的,不在数组中
        if (sum < k) {
            return Integer.MIN_VALUE;
        }

        // 跳过第一个if,说明,元素一定在数组中
        // 如果加入的个数大于k,说明需要的元素已经在队列中了,只需取出来即可
        if (count >= k) {
            for (int i = 1; i < k; i++) {
                pq.poll();
            }
            return pq.poll()[2];
        }

        // 跳过第二个if,说明需要的元素在数组中,但是还没有加入到优先队列中
        // 此时,一定要从优先队列中出去的元素来进行记录,出去3个,说明找到了前3个最小的元素,count用来记录出去的个数
        count = 0;
        int[] head = null;
        while (count != k) {
            head = pq.poll();
            count++;
            int row = head[0];
            int col = head[1];
            if (col + 1 < nums[row].length) {
                int[] next = new int[]{row, col + 1, nums[row][col + 1]};
                pq.offer(next);
            }
        }

        return head[2];
    }

results matching ""

    No results matching ""