Поиск лучших K элементов в массиве Java

1. Обзор

В этом руководстве мы реализуем различные решения проблемы поиска k самых больших элементов в массиве с помощью Java. Для описания временной сложности мы будем использовать нотацию Big-O .

2. Решение грубой силы

Грубое решение этой проблемы состоит в том, чтобы перебрать заданный массив k раз . На каждой итерации мы будем находить наибольшее значение . Затем мы удалим это значение из массива и поместим в выходной список:

public List findTopK(List input, int k) {
    List array = new ArrayList<>(input);
    List topKList = new ArrayList<>();

    for (int i = 0; i < k; i++) {
        int maxIndex = 0;

        for (int j = 1; j < array.size(); j++) {
            if (array.get(j) > array.get(maxIndex)) {
                maxIndex = j;
            }
        }

        topKList.add(array.remove(maxIndex));
    }

    return topKList;
}

Если мы предположим, что n — это размер данного массива, временная сложность этого решения равна O(n * k) . Кроме того, это самое неэффективное решение.

3. Подход с коллекциями Java

Однако существуют более эффективные решения этой проблемы. В этом разделе мы объясним два из них с использованием коллекций Java.

3.1. Набор деревьев

TreeSet имеет структуру данных Red-Black Tree в качестве основы. В результате добавление значения в этот набор стоит O(log n) . TreeSet — это отсортированная коллекция. Следовательно, мы можем поместить все значения в TreeSet и извлечь первые k из них: ** ** **** ** **

public List<Integer> findTopK(List<Integer> input, int k) {
    Set<Integer> sortedSet = new TreeSet<>(Comparator.reverseOrder());
    sortedSet.addAll(input);

    return sortedSet.stream().limit(k).collect(Collectors.toList());
}

Временная сложность этого решения составляет O(n * log n) . Прежде всего, предполагается, что это более эффективно, чем метод грубой силы, если k ≥ log n .

Важно помнить, что TreeSet не содержит дубликатов. В результате решение работает только для входного массива с различными значениями.

3.2. PriorityQueue

PriorityQueue — это структура данных Heap в Java. С его помощью мы можем добиться решения O(n * log k) . Более того, это будет более быстрое решение, чем предыдущее. Из-за указанной проблемы k всегда меньше размера массива. Итак, это означает, что O(n * log k) ≤ O(n * log n). ** ** ``

Алгоритм проходит один раз по заданному массиву . На каждой итерации мы будем добавлять новый элемент в кучу. Кроме того, мы сохраним размер кучи меньше или равным k . Итак, нам придется удалить лишние элементы из кучи и добавить новые. В результате после перебора массива куча будет содержать k самых больших значений:

public List<Integer> findTopK(List<Integer> input, int k) {
    PriorityQueue<Integer> maxHeap = new PriorityQueue<>();

    input.forEach(number -> {
        maxHeap.add(number);

        if (maxHeap.size() > k) {
            maxHeap.poll();
        }
    });

    List<Integer> topKList = new ArrayList<>(maxHeap);
    Collections.reverse(topKList);

    return topKList;
}

4. Алгоритм выбора

Существует множество подходов к решению данной проблемы. И, хотя это выходит за рамки данного руководства, использование подхода алгоритма выбора будет лучшим, поскольку он дает линейную временную сложность.

5. Вывод

В этом руководстве мы описали несколько решений для поиска k самых больших элементов в массиве.

Как обычно, код примера доступен на GitHub .