321. 拼接最大数

题目描述

输入两个整数数组 nums1 和 nums2，长度分别为 m 和 n，元素都是 0 到 9 的数字，分别表示两个数从高到低的各位数字。另外输入一个整数 k，满足 k <= m + n。

从两个数组中一共选出 k 个数字，拼成一个 k 位的数。来自同一个数组的数字必须保持它们在原数组里的先后顺序，两个数组的数字之间可以任意穿插。

让拼出来的数尽可能大，返回它的各位数字组成的长度为 k 的数组。

示例 1

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

解释：从 nums1 取 6,5，从 nums2 取 9,8,3，合并成 [9,8,6,5,3]。

示例 2

输入：nums1 = [6,7], nums2 = [6,0,4], k = 5
输出：[6,7,6,0,4]

解释：两个数组的数字全都要用上，交错排列后最大的是 [6,7,6,0,4]。

示例 3

输入：nums1 = [3,9], nums2 = [8,9], k = 3
输出：[9,8,9]

解释：从 nums2 取 8,9，从 nums1 取 9，最大是 [9,8,9]。

数据范围

m == nums1.length
n == nums2.length
1 <= m, n <= 500
0 <= nums1[i], nums2[i] <= 9
1 <= k <= m + n
nums1 和 nums2 都没有前导 0。

解题思路

把问题拆成三层。第一层：枚举从 nums1 取 i 个、从 nums2 取 k-i 个。第二层：在单个数组里取固定 t 个数字使其最大，这是单调栈问题，可以删掉 len-t 个数字：遇到更大的数字时，只要还有删除额度，就弹出比它小的栈顶。第三层：把两个子序列合并成最大的序列，每一步比较两边“剩余的整个后缀”的字典序，取较大一边的首位（只比首位在相等时会出错）。所有拆分的结果中取字典序最大的一个。

为什么正确

pick 过程中，栈加上剩余元素始终至少有 t 个（drop 只会在仍有富余时减少），并且在可用的删除额度内，栈是单调不增的；merge 时每次取的都是“后缀字典序更大”的那一边的首位，所以合并结果保留了两边能提供的最大前缀。

复杂度分析

时间 O(k·(m+n)²)：共 O(k) 种拆分，每种 pick 是 O(m+n)，merge 每步比较后缀最坏 O(m+n)，共 O((m+n)²)；空间 O(m+n)。

易错点

合并时不能只比较两个首位：首位相等时要继续比较后面，否则比如 [6,7] 与 [6,0,4] 会先取错 6，得到 [6,6,7,0,4] 而不是 [6,7,6,0,4]。另外枚举 i 的范围是 max(0, k-n) 到 min(k, m)，越界会让 pick 拿到负数个数字。

参考代码

Python、Java、C++、Go、JavaScript、TypeScript 六种语言，逻辑一致。

class Solution:
    def maxNumber(self, nums1: List[int], nums2: List[int], k: int) -> List[int]:
        def pick(nums, t):
            drop = len(nums) - t
            stack = []
            for x in nums:
                while drop and stack and stack[-1] < x:
                    stack.pop()
                    drop -= 1
                stack.append(x)
            return stack[:t]

        def merge(a, b):
            res = []
            while a or b:
                if a > b:
                    res.append(a.pop(0))
                else:
                    res.append(b.pop(0))
            return res

        best = []
        for i in range(max(0, k - len(nums2)), min(k, len(nums1)) + 1):
            cand = merge(pick(nums1, i), pick(nums2, k - i))
            if cand > best:
                best = cand
        return best

class Solution {
    private int[] pick(int[] nums, int t) {
        int drop = nums.length - t;
        int[] stack = new int[nums.length];
        int top = 0;
        for (int x : nums) {
            while (drop > 0 && top > 0 && stack[top - 1] < x) {
                top--;
                drop--;
            }
            stack[top++] = x;
        }
        return Arrays.copyOf(stack, t);
    }

    private boolean greater(int[] a, int i, int[] b, int j) {
        while (i < a.length && j < b.length && a[i] == b[j]) {
            i++;
            j++;
        }
        return j == b.length || (i < a.length && a[i] > b[j]);
    }

    private int[] merge(int[] a, int[] b) {
        int[] res = new int[a.length + b.length];
        int i = 0, j = 0, r = 0;
        while (i < a.length || j < b.length) {
            if (greater(a, i, b, j)) {
                res[r++] = a[i++];
            } else {
                res[r++] = b[j++];
            }
        }
        return res;
    }

    public int[] maxNumber(int[] nums1, int[] nums2, int k) {
        int[] best = new int[k];
        for (int i = Math.max(0, k - nums2.length); i <= Math.min(k, nums1.length); i++) {
            int[] cand = merge(pick(nums1, i), pick(nums2, k - i));
            if (greater(cand, 0, best, 0)) {
                best = cand;
            }
        }
        return best;
    }
}

class Solution {
    vector<int> pick(vector<int>& nums, int t) {
        int drop = nums.size() - t;
        vector<int> stk;
        for (int x : nums) {
            while (drop && !stk.empty() && stk.back() < x) {
                stk.pop_back();
                drop--;
            }
            stk.push_back(x);
        }
        stk.resize(t);
        return stk;
    }
    vector<int> merge(vector<int>& a, vector<int>& b) {
        vector<int> res;
        int i = 0, j = 0;
        while (i < (int)a.size() || j < (int)b.size()) {
            if (lexicographical_compare(b.begin() + j, b.end(), a.begin() + i, a.end())) {
                res.push_back(a[i++]);
            } else {
                res.push_back(b[j++]);
            }
        }
        return res;
    }
public:
    vector<int> maxNumber(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<int> best;
        int m = nums1.size(), n = nums2.size();
        for (int i = max(0, k - n); i <= min(k, m); i++) {
            vector<int> a = pick(nums1, i), b = pick(nums2, k - i);
            vector<int> cand = merge(a, b);
            if (cand > best) {
                best = cand;
            }
        }
        return best;
    }
};

func pick(nums []int, t int) []int {
	drop := len(nums) - t
	stack := []int{}
	for _, x := range nums {
		for drop > 0 && len(stack) > 0 && stack[len(stack)-1] < x {
			stack = stack[:len(stack)-1]
			drop--
		}
		stack = append(stack, x)
	}
	return stack[:t]
}

func greater(a, b []int) bool {
	i := 0
	for i < len(a) && i < len(b) && a[i] == b[i] {
		i++
	}
	return i < len(a) && (i == len(b) || a[i] > b[i])
}

func merge(a, b []int) []int {
	res := []int{}
	for len(a) > 0 || len(b) > 0 {
		if greater(a, b) {
			res = append(res, a[0])
			a = a[1:]
		} else {
			res = append(res, b[0])
			b = b[1:]
		}
	}
	return res
}

func maxNumber(nums1 []int, nums2 []int, k int) []int {
	best := []int{}
	lo, hi := k-len(nums2), k
	if lo < 0 {
		lo = 0
	}
	if hi > len(nums1) {
		hi = len(nums1)
	}
	for i := lo; i <= hi; i++ {
		cand := merge(pick(nums1, i), pick(nums2, k-i))
		if greater(cand, best) {
			best = cand
		}
	}
	return best
}

var maxNumber = function(nums1, nums2, k) {
    const pick = (nums, t) => {
        let drop = nums.length - t;
        const stack = [];
        for (const x of nums) {
            while (drop > 0 && stack.length && stack[stack.length - 1] < x) {
                stack.pop();
                drop--;
            }
            stack.push(x);
        }
        return stack.slice(0, t);
    };
    const greater = (a, i, b, j) => {
        while (i < a.length && j < b.length && a[i] === b[j]) {
            i++;
            j++;
        }
        return i < a.length && (j === b.length || a[i] > b[j]);
    };
    const merge = (a, b) => {
        const res = [];
        let i = 0, j = 0;
        while (i < a.length || j < b.length) {
            if (greater(a, i, b, j)) {
                res.push(a[i++]);
            } else {
                res.push(b[j++]);
            }
        }
        return res;
    };
    let best = [];
    for (let i = Math.max(0, k - nums2.length); i <= Math.min(k, nums1.length); i++) {
        const cand = merge(pick(nums1, i), pick(nums2, k - i));
        if (greater(cand, 0, best, 0)) {
            best = cand;
        }
    }
    return best;
};

function maxNumber(nums1: number[], nums2: number[], k: number): number[] {
    const pick = (nums: number[], t: number): number[] => {
        let drop = nums.length - t;
        const stack: number[] = [];
        for (const x of nums) {
            while (drop > 0 && stack.length && stack[stack.length - 1] < x) {
                stack.pop();
                drop--;
            }
            stack.push(x);
        }
        return stack.slice(0, t);
    };
    const greater = (a: number[], i: number, b: number[], j: number): boolean => {
        while (i < a.length && j < b.length && a[i] === b[j]) {
            i++;
            j++;
        }
        return i < a.length && (j === b.length || a[i] > b[j]);
    };
    const merge = (a: number[], b: number[]): number[] => {
        const res: number[] = [];
        let i = 0, j = 0;
        while (i < a.length || j < b.length) {
            if (greater(a, i, b, j)) {
                res.push(a[i++]);
            } else {
                res.push(b[j++]);
            }
        }
        return res;
    };
    let best: number[] = [];
    for (let i = Math.max(0, k - nums2.length); i <= Math.min(k, nums1.length); i++) {
        const cand = merge(pick(nums1, i), pick(nums2, k - i));
        if (greater(cand, 0, best, 0)) {
            best = cand;
        }
    }
    return best;
}

载入编辑器，自己跑一遍 →

栈与堆专题的其他题目

查看栈与堆全部题目 →