找到第 m 个最大的元素

Question

给定一个

n

整数数组，数组中的值是任意顺序的

1 to n

。

给出另一个整数

m

作为输入。

现在从数组（上述数组的子数组）中选取

m

项，并从此子数组中找到

mth largest element

。

接下来从数组中选择

m+1

项并从此子数组中找到

mth largest element

依此类推..直到子数组包含所有

n

项。

示例：

n = 4
arr = [4,2,1,3]
m = 2

结果：

[2,2,3]

说明：

a) m items = [4,2] , mth largest = 2nd largest = 2
b) m+1 items = [4,2,1], 2nd largest = 2
c) m+2 items = [4,2,1,3], 2nd largest = 3

so result is [2,2,3]

这是我的代码：

static List<Integer> solve(List<Integer> arr, int m) {
   List<Integer> result = new ArrayList<>();
   int n = arr.size();
   for(int i=m; i<=n; i++) {
    List<Integer> sub = arr.subList(0, i);
    Collections.sort(sub, Collections.reverseOrder());
    result.add(sub.get(m-1));
   }
   return result;
}

我想优化它，使其运行时间更少，解决这个问题的正确方法是什么？

Answer 1

您可以通过维护一个包含运行前缀中的 m 最大值的二进制堆来在 O(N log N) 时间内解决此问题。第 m 个最大值就是这些 m 值中最小的一个。

这是一个用 Python 实现的示例来演示该算法：

import heapq

n = 4
arr = [4,2,1,3]
m = 2

heap = arr[0:m]
heapq.heapify(heap)
res = []
for i in range(m, n):
  # Invariant: heap contains the m largest values of arr[0:i] in min-heap order.
  # Therefore, heap[0] (the smallest of the m largest values) is the m-largest value in arr[0:m].
  res.append(heap[0])
  # Restore the invariant by adding arr[i] and removing the smallest value from heap.
  heapq.heappushpop(heap, arr[i])
# At this point, heap contains the m largest values of arr. Add the final value to the result.
res.append(heap[0])

print(res)  # prints [2, 2, 3]

您可以使用 PriorityQueue 在 Java 中实现此功能。您将使用接受集合的构造函数，而不是

heapify

，而不是

heap[0]

，您可以调用

peek()

，而不是

heappushpop

，您可以调用

add()

，然后调用

poll()

。