求修改O(log(m*n))复杂度双有序数组第k小算法为第k大算法

我有一段C++代码，能以O(log(m*n))的时间复杂度找出两个有序数组中的第k小元素。我知道存在O(log(min(m,n)))的最优解法，但目前暂不关注。

我尝试修改这段代码使其返回第k大元素，但调整比较逻辑后始终无法正常运行。我明白只需微调部分比较逻辑，但一直没成功，求帮助。

原第k小元素代码 #

#include <iostream>
#include <vector>
#include <string>
#include <unordered_set>
#include <unordered_map>

using namespace std;

// speed: O(log( m * n) ))
// space: O(1)
int solveKthSmallest(vector<int>& A, vector<int>& B, int k, int aStart, int aEnd, int bStart, int bEnd)
{
    std::cout << "solveKthSmallest: (" << k << "): [" << aStart << "," << aEnd << "] [" << bStart << "," << bEnd << "\n";

    // 若其中一个数组的当前段为空，直接返回另一个数组中对应元素
    if (aEnd < aStart) {
        return B[k - aStart];
    }
    if (bEnd < bStart) {
        return A[k - bStart];
    }

    // 获取当前段的中间索引和对应值
    int aIndex = (aStart + aEnd) / 2;
    int bIndex = (bStart + bEnd) / 2;
    int aValue = A[aIndex];
    int bValue = B[bIndex];

    // 若k在合并数组的右半部分，移除较小的左半部分
    if (aIndex + bIndex < k) { 
        if (aValue > bValue) {
            return solveKthSmallest(A, B, k, aStart, aEnd, bIndex + 1, bEnd);
        } else {
            return solveKthSmallest(A, B, k, aIndex + 1, aEnd, bStart, bEnd);
        }
    } 
    // 否则，移除较大的右半部分
    else { 
        if (aValue > bValue) {
            return solveKthSmallest(A, B, k, aStart, aIndex - 1, bStart, bEnd);
        } else {
            return solveKthSmallest(A, B, k, aStart, aEnd, bStart, bIndex - 1);       
        }
    }
    return -1;
}

double findKthSmallest(vector<int>& nums1, vector<int>& nums2, int kIdx) {
    return solveKthSmallest(nums1, nums2, kIdx, 0, nums1.size() - 1, 0, nums2.size() - 1);
}

void t1()
{
    vector nums1{ 1, 4, 5, 8, 9 };
    vector nums2{ 2, 3, 6, 7 };

    cout << findKthSmallest(nums1, nums2, 3) << "\n";
}

int main()
{
    t1();

    return 0;
}

已知的间接实现方式 #

我知道可以通过转换索引间接实现，代码如下：

double findKthLargest(vector<int>& nums1, vector<int>& nums2, int kIdx) {
    int revIdx = nums1.size() + nums2.size() - 1 - kIdx;
    cout << "revIdx: " << revIdx << "\n";

    return solveKthSmallest(nums1, nums2, revIdx, 0, nums1.size() - 1, 0, nums2.size() - 1);
}

但我希望复制并修改solveKthSmallest函数为solveKthLargest，保留原算法逻辑，直接通过以下方式调用：

return solveKthLargest(nums1, nums2, kIdx, 0, nums1.size() - 1, 0, nums2.size() - 1);

修改后的直接实现方案 #

完整的solveKthLargest代码 #

#include <iostream>
#include <vector>

using namespace std;

// 时间复杂度O(log(m*n))，空间复杂度O(1)
int solveKthLargest(vector<int>& A, vector<int>& B, int k, int aStart, int aEnd, int bStart, int bEnd)
{
    std::cout << "solveKthLargest: (" << k << "): [" << aStart << "," << aEnd << "] [" << bStart << "," << bEnd << "\n";

    // 若其中一个数组的当前段为空，直接返回另一个数组中对应第k大的元素
    if (aEnd < aStart) {
        return B[bEnd - k + aStart];
    }
    if (bEnd < bStart) {
        return A[aEnd - k + bStart];
    }

    // 获取当前段的中间索引和对应值
    int aIndex = (aStart + aEnd) / 2;
    int bIndex = (bStart + bEnd) / 2;
    int aValue = A[aIndex];
    int bValue = B[bIndex];

    // 计算当前中间位置右侧的元素总数（包括中间元素本身）
    int rightCount = (aEnd - aIndex) + (bEnd - bIndex) + 2;

    // 如果k在右侧区域，保留较大的右侧部分，移除较小的左侧部分
    if (rightCount > k) { 
        if (aValue > bValue) {
            // A的中间元素更大，保留A的右侧，移除B的右侧
            return solveKthLargest(A, B, k, aIndex + 1, aEnd, bStart, bEnd);
        } else {
            // B的中间元素更大，保留B的右侧，移除A的右侧
            return solveKthLargest(A, B, k, aStart, aEnd, bIndex + 1, bEnd);
        }
    } 
    // 否则，移除较大的右侧部分，在左侧区域继续查找（k要减去被移除的元素数量）
    else { 
        if (aValue > bValue) {
            // A的中间元素更大，移除A的右侧，剩余查找k减去B右侧的元素数
            return solveKthLargest(A, B, k - (bEnd - bIndex + 1), aStart, aIndex - 1, bStart, bEnd);
        } else {
            // B的中间元素更大，移除B的右侧，剩余查找k减去A右侧的元素数
            return solveKthLargest(A, B, k - (aEnd - aIndex + 1), aStart, aEnd, bStart, bIndex - 1);       
        }
    }
    return -1;
}

double findKthLargest(vector<int>& nums1, vector<int>& nums2, int kIdx) {
    return solveKthLargest(nums1, nums2, kIdx, 0, nums1.size() - 1, 0, nums2.size() - 1);
}

// 测试用例
void t2()
{
    vector nums1{ 1, 4, 5, 8, 9 };
    vector nums2{ 2, 3, 6, 7 };

    // 合并后数组：[1,2,3,4,5,6,7,8,9]，第1大是9（kIdx=0），第2大是8（kIdx=1）
    cout << findKthLargest(nums1, nums2, 1) << "\n"; // 输出8
}

int main()
{
    t2();
    return 0;
}

关键修改点说明 #

1. 边界索引计算：当一个数组为空时，第k大元素的索引需要从数组末尾往前推算，而非原算法的从开头往后数
2. 右侧元素计数逻辑：计算两个数组当前中间位置到末尾的元素总数，用来判断目标k是否落在这个更大元素聚集的右侧区域
3. 递归分支反转：
   • 当k在右侧区域时，保留较大元素所在数组的右侧部分，移除较小元素所在数组的右侧
   • 当k不在右侧区域时，移除较大元素所在数组的右侧，并将k减去被移除区域的元素数量，在剩余的左侧区域继续查找

内容的提问来源于stack exchange，提问作者Mikhail Kalashnikov