Поиск k-ой порядковой статистики в двух массивах

Задача:

Даны два отсортированных массива [math]A[/math] и [math]B[/math] размерами [math]n[/math] и [math]m[/math] соответственно. Требуется найти [math]k[/math]-ый порядковый элемент после их слияния. Будем считать, что все элементы в массивах различны и нумеруются с нуля.

Варианты решения

Наивное решение

Сольем два массива и просто возьмем элемент с индексом [math]k - 1[/math]. Слияние будет выполнено за время [math]O(n + m)[/math], к тому же этот алгоритм использует [math]O(n + m)[/math] дополнительной памяти.

Чуть менее наивное решение

Будем использовать два указателя, с помощью которых сможем обойти массивы, не сливая их. Поставим указатели на начало каждого из массивов. Будем увеличивать на единицу тот из них, который указывает на меньший элемент. После [math](k - 1)[/math]-ой итерации сравним элементы, на которых стоят указатели. Меньший из них и будет ответом. Таким образом, мы получим [math]k[/math]-ый элемент за [math]O(k)[/math] шагов.

Еще одно решение

В первом массиве выберем элемент c индексом [math]i = \dfrac{n}{2}[/math] и бинарным поиском найдем во втором массиве позицию [math]j[/math], на которой стоит наибольший элемент, меньший [math]a[i][/math]. Если [math]i + j = k - 2[/math], то мы нашли [math]k[/math]-ую порядковую статистику — это элемент [math]a[i][/math]. Иначе, если [math]i + j \gt k - 2[/math], то далее тем же способом ищем в массиве [math]A[/math] в диапазоне индексов [math][0, i - 1][/math], а если [math]i + j \lt k - 2[/math], то в диапазоне индексов [math][i + 1, n - 1][/math]. Решая задачу таким способом, мы получим асимптотику [math]O(\log(n) \cdot \log(m))[/math].

Совсем не наивное решение

Приведём теперь решение, работающее за время [math]O(\log(\min(n, m)))[/math].

Для начала рассмотрим следующую ситуацию: пусть у нас есть элемент [math]a[i][/math] из массива [math]A[/math] и элемент [math]b[j][/math] из массива [math]B[/math] и они связаны неравенством [math]b[j - 1] \lt a[i] \lt b[j][/math]. Тогда [math]a[i][/math] есть [math](j + i + 1)[/math]-ый порядковый элемент после слияния массивов. Это объясняется тем, что до [math]a[i][/math]-ого элемента идут [math]j[/math] элементов из массива [math]B[/math], [math](i+1)[/math] элементов из массива [math]A[/math] (включая сам элемент [math]a[i][/math]). В итоге получаем [math]j + i + 1[/math]. Принимая это во внимание, будем выбирать [math]i[/math] и [math]j[/math] таким образом, чтобы [math]j + i + 1 = k[/math].

Подведем промежуточный итог:

1. Инвариант [math]j + i = k - 1[/math]
2. Если [math]b[j - 1] \lt a[i] \lt b[j][/math], то [math]a[i][/math] и есть [math]k[/math]-ая порядковая статистика
3. Если [math]a[i - 1] \lt b[j] \lt a[i][/math], то [math]b[j][/math] и есть [math]k[/math]-ая порядковая статистика

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

Будем использовать [math]i[/math] и [math]j[/math] как опорные точки для разделения массивов. Заметим, что если [math]a[i] \lt b[j][/math], то [math]a[i] \lt b[j - 1][/math] (иначе второе условие бы выполнялось). В таком случае на месте [math]i[/math]-го элемента может стоять максимум [math]i + (j - 2) + 2 = (i + j)[/math]-ый порядковый элемент после слияния массивов (так произойдет в случае, когда [math]a[i] \gt b[j - 2][/math]), а значит элемент с номером [math]i[/math] и все до него в массиве [math]A[/math] никогда не будут [math]k[/math]-ой порядковой статистикой. Аналогично элемент с индексом [math]j[/math] и все элементы, стоящие после него, в массиве [math]B[/math] никогда не будут ответом, так как после слияния на позиции [math]j[/math] будет стоять [math](i + j + 2)[/math]-ой порядковый элемент, порядковые номера остальных же будут еще больше. Таким образом, далее мы можем продолжать поиск в массиве [math]A[/math] только в диапазоне индексов [math][i + 1, n - 1][/math], а в массиве [math]B[/math] — [math][0, j - 1][/math]. По аналогии, если [math]b[j] \lt a[i][/math], то [math]b[j] \lt a[i - 1][/math] (иначе выполнялось бы третье условие). Аналогичными рассуждениями приходим к тому, что в таком случае дальнейший поиск нужно осуществлять в массиве [math]A[/math] в диапазоне [math][0, i - 1][/math], в массиве [math]B[/math] — [math][j + 1, m - 1][/math].

Стоит отметить, что нам не нужно рассматривать элементы, стоящие и в том, и в другом массивах на позициях от [math]k[/math]-ой до конца (если такие есть), так как они тоже никогда не будут ответом. Поэтому первый раз запускаем нашу функцию от параметров [math]\mathtt{findKthOrderStatistic}(A, \min(n, k), B, \min(m, k), k)[/math].

int findKthOrderStatistic(int* A, int n, int* B, int m, int k): 
  if n == 1 // в этом случае можно сразу дать ответ 
    if B[k - 1] < A[0]
      return B[k - 1]
    else if  A[0] < B[k - 2]
      return B[k - 2]
    else
      return A[0]
  if m == 1 // симметричен случаю с n = 1 
      return findKthOrderStatistic(B, m, A, n, k)
  int i = n / 2
  int j = (k - 1) - i // j > 0, так как i <= (k / 2) 
  if j >= m
    return findKthOrderStatistic(A + i + 1, n - i - 1, B, m, k - i - 1)
  // чтобы сохранить инвариант, сделаем A[-1] = -INF и B[-1] = -INF 
  int AiLeft = ((i == 0) ? INT_MIN : A[i - 1])
  int BjLeft = ((j == 0) ? INT_MIN : B[j - 1])
  if BjLeft < A[i] and A[i] < B[j]
    return A[i]
  else if AiLeft < B[j] and B[j] < A[i]
    return B[j]
  if A[i] < B[j]
    return findKthOrderStatistic(A + i + 1, n - i - 1, B, j, k - i - 1)
  else
    return findKthOrderStatistic(A, i, B + j + 1, m - j - 1, k - j - 1)

Чтобы алгоритм работал за [math]O(\log(\min(n, m)))[/math], будем передавать первым массивом в функцию тот, длина которого меньше. Тогда длина рассматриваемой области первого массива на каждой итерации уменьшается в два раза. Как только она станет равна единице, то за несколько сравнений легко получить ответ.

См. также

• Сортировка слиянием
• Быстрая сортировка
• Поиск k-ой порядковой статистики за линейное время

Источники информации

• LeetCode — Find the k-th Smallest Element in the Union of Two Sorted Arrays
• Blogspot — A Cute Algorithm