来自两个给定数组的子序列的最大K位数字

问题描述

给定两个非负整数数组 nums1 和 nums2 ，长度分别为 n1 和 n2 。假设两个数组都不为空且长度都不为零。

现在你需要从两个数组中选出 k 个数字（k <= n1 + n2），构成数字序列 S ，使得该数字序列的从左往右的数值最大。

请你在可行的数字序列 S 上面，选取字典序最大的一种方案。

示例

输入：
nums1 = [3, 4, 6, 5]
nums2 = [9, 1, 2, 5, 8, 3]
k = 5

输出：
[9, 8, 6, 5, 3]

解法

这个问题可以转化为一个单调栈问题。

考虑从 nums1 中选择 i 个数字，从 nums2 中选择 j 个数字，使得 i+j=k，且这个选择方式可以得到最终的数值最大的 S。

假设我们已经选择了 nums1 中的前 i 个数字，那么此时要选择 nums2 中的前 j 个数字，使得选择的数字能够和 nums1 中已经选择的数字合成最终的数字最大。而这需要一个单调栈来保证。

具体来说，我们可以定义一个单调栈，用于保存 nums2 中前 j 个数字中可用的数字。从第一个到第 j 个数字遍历 nums2，将数字加入单调栈中。每当单调栈的长度超过了 j-i 个数字时，就弹出栈顶元素。最终栈内剩下的数字就是剩余可用的数字。

然后我们将 nums1[i] 与单调栈中的数字依次比较，选出最大的数字，就是合并 nums1 前 i 个数字和 nums2 中可用的数字后得到的最大数值。

代码

class Solution:
    def maxNumber(self, nums1: List[int], nums2: List[int], k: int) -> List[int]:
        m, n = len(nums1), len(nums2)
        ans = [0] * k

        start, end = max(0, k - n), min(k, m)
        for i in range(start, end + 1):
            sub1 = self.findMaxSubsequence(nums1, i)
            sub2 = self.findMaxSubsequence(nums2, k - i)
            curMaxSubsequence = self.merge(sub1, sub2)
            if self.compare(curMaxSubsequence, 0, ans, 0) > 0:
                ans = curMaxSubsequence
        
        return ans

    def findMaxSubsequence(self, nums: List[int], k: int) -> List[int]:
        stack = []
        drop = len(nums) - k
        for num in nums:
            while drop and stack and stack[-1] < num:
                stack.pop()
                drop -= 1
            stack.append(num)
        return stack[:k]

    def merge(self, nums1: List[int], nums2: List[int]) -> List[int]:
        ans = [0] * (len(nums1) + len(nums2))
        i, j, k = 0, 0, 0
        while i < len(nums1) and j < len(nums2):
            if self.compare(nums1, i, nums2, j) > 0:
                ans[k] = nums1[i]
                i += 1
            else:
                ans[k] = nums2[j]
                j += 1
            k += 1
        
        while i < len(nums1):
            ans[k] = nums1[i]
            i += 1
            k += 1
        
        while j < len(nums2):
            ans[k] = nums2[j]
            j += 1
            k += 1
        
        return ans

    def compare(self, nums1: List[int], i: int, nums2: List[int], j: int) -> int:
        while i < len(nums1) and j < len(nums2):
            diff = nums1[i] - nums2[j]
            if diff != 0:
                return diff
            i += 1
            j += 1
        return (len(nums1) - i) - (len(nums2) - j)

代码片段来源：力扣（LeetCode）

时间复杂度

• 遍历数组 nums1，并将 nums1 分别合并到 nums2 的前 i 个数字和后 n-i 个数字中，时间复杂度为 O(k^3)。

• 对于每个 i，需要将 nums1 中的前 i 个数字与 nums2 中的前 k-i 个数字进行合并，时间复杂度为 O(k^2)。

• 对于每个 i，需要使用单调栈从 nums2 中选出 k-i 个数字，时间复杂度也为 O(k^2)。

综上，总时间复杂度为 O(k^3)。