K-ésimo menor elemento em uma matriz classificada

Question

Portanto, estou trabalhando em uma questão Leetcode e meu código funciona para alguns casos, mas falha em alguns casos.

Aqui está a questão:

Dada uma n x nmatriz onde cada uma das linhas e colunas são classificadas em ordem crescente, encontre o k-ésimo menor elemento na matriz.

Observe que é o k-ésimo menor elemento na ordem de classificação, não o k-ésimo elemento distinto.

Exemplo:

matrix = [[1, 5, 9], [10, 11, 13], [12, 13, 15]]
k = 8

retorno: 13

Minha abordagem é usar um minHeap, mesmo que afirme que o array está classificado, eu ainda preciso ter certeza de que o classifiquei do menor para o maior valor.

Aqui está o meu código:

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int row = matrix.length;
        int col = matrix[0].length;
        int result = 0;

        HashMap<Integer, Integer> map = new HashMap<>();
        //populate HashMap 
        for (int i = 0; i < row; i++) {
            for (int j = 0; j < col; j++) {
                map.put(matrix[i][j],
                        map.getOrDefault(matrix[i][j], 0) + 1);
            }
        }

        PriorityQueue<Map.Entry<Integer, Integer>> pq =
                new PriorityQueue<>((n1, n2) -> n1.getValue() - n2.getValue());

        pq.addAll(map.entrySet());

        for (int i = 0; i < k && !(pq.isEmpty()); i++) {
            result = pq.poll().getKey();
        }

        return result;
    }
}

Aqui estão minhas informações:

Input 1: [[1,5,9],[10,11,13],[12,13,15]] 
k = 8 

Input 2: [[1,2],[1,3]] 
k = 1 

Aqui estão os resultados:

Output 1: 13
Output 2: 2

Observe que o código funciona bem para minha primeira entrada, onde o 8º menor elemento no array 2d é 13, mas para a segunda entrada, o código está retornando 2 como meu primeiro menor elemento em vez de retornar 1.

Alguém pode me ajudar a consertar o código? Peço que você não implemente alguma técnica sofisticada de classificação abreviada, por exemplo, Arrays.sort ... não é ideal para mim, pois estou tentando aprender como implementar heaps. Muito obrigado :)

Answer 1

Para resolver este problema, também podemos pesquisar binário (um pouco mais eficiente):

public class Solution {
    public static final int kthSmallest(final int[][] matrix, final int k) {
        int lo = matrix[0][0];
        int hi = matrix[matrix.length - 1][matrix[0].length - 1] + 1;

        while (lo < hi) {
            final int mid = lo + (hi - lo) / 2;
            int count = 0;
            int col = matrix[0].length - 1;

            for (int row = 0; row < matrix.length; ++row) {
                while (col >= 0 && matrix[row][col] > mid) {
                    col--;
                }
                count += (col + 1);
            }
            if (count < k) {
                lo = mid + 1;
            } else {
                hi = mid;
            }
        }
        return lo;
    }
}

Answer 2

Você pode usar o flatMapToIntmétodo para iterar sobre as linhas de uma matriz 2d em um fluxo:

public static void main(String[] args) {
    int[][] matrix1 = {{1, 5, 9}, {10, 11, 13}, {12, 13, 15}};
    int[][] matrix2 = {{1, 2}, {1, 3}};

    System.out.println(kthSmallest(matrix1, 8)); // 13
    System.out.println(kthSmallest(matrix2, 1)); // 1
}

public static int kthSmallest(int[][] matrix, int k) {
    return Arrays.stream(matrix)
            .flatMapToInt(Arrays::stream)
            .skip(k - 1)
            .findFirst()
            .getAsInt();
}