亚马逊面试题：求整数数组元素最少且和大于剩余部分的子集

题目要求 #

给定整数数组arr，编写代码返回符合以下规则的子集A：

• 子集A与剩余元素组成的子集B交集为空
• 子集A与子集B的并集为原始数组
• 子集A的元素数量为所有符合条件子集中的最小值
• 子集A的元素之和严格大于子集B的元素之和
• 满足元素数量最少的前提下，优先选择元素和最大的子集

以数组3,7,5,6,2为例，元素数为2的符合和要求的子集有两个：

• 子集5,7，和为12，大于剩余子集3,6,2的和11
• 子集6,7，和为13，大于剩余子集3,5,2的和10，为本题正确输出

原有实现仅通过13个测试用例中的2个，无法满足全部规则要求。

原有代码问题分析 #

原有实现代码如下：

public static List<int> minimalHeaviestSetA(List<int> arr)
{
    List<int> sol = new List<int>();List<int> sec= new List<int>();
    //try finding the two largest numbers in the array and add
    int max=0;
    foreach(int u in arr){
        if(u>max){
            max=u;
        }
    }
    //add the maximum element to the list
    
    //initialize the sum
    int sum=0;
    //remove the maximum element from the list
    arr.Remove(max);
    foreach(int p in arr){
        if((p+max)>sum){
            sum=p+max;
            sec.Clear();
            sec.Add(p);
        }
        
    }
    sol.Add(sec[0]);
    sol.Add(max);
    return sol;
    
}

核心问题有3个：

1. 硬编码假设结果固定为2个元素，完全没有覆盖单元素满足要求、需要3个及以上元素才能满足和要求的场景
2. 最大值计算逻辑存在缺陷，遇到包含负数的数组会直接失效，且未对数组做全局排序，无法按优先级选取最大元素
3. 没有计算数组全局总和，无法动态判断当前选中的元素和是否已经超过剩余元素和，同时用int存储和容易出现整数溢出问题

正确解题思路 #

要满足「元素数最少、同元素数下和最大」的要求，最优策略是贪心：

1. 先计算数组所有元素的总和，判断阈值为总和的1/2，只要选中元素的和严格超过这个阈值就满足和的要求
2. 将数组按降序排序，要凑最少元素超过阈值，必然优先选数值最大的元素
3. 从排序后的数组头部开始依次累加元素，直到当前累加和严格大于剩余元素的和就停止选取
4. 此时选中的元素就是符合要求的子集，按题目要求升序排列后返回即可

正确实现代码 #

public static List<int> minimalHeaviestSetA(List<int> arr)
{
    // 用long存储总和避免整数溢出
    long totalSum = 0;
    foreach (int num in arr)
    {
        totalSum += num;
    }
    // 数组降序排序，优先取大值
    arr.Sort((a, b) => b.CompareTo(a));
    List<int> result = new List<int>();
    long currentSum = 0;
    foreach (int num in arr)
    {
        currentSum += num;
        result.Add(num);
        // 满足当前和严格大于剩余和，停止选取
        if (currentSum > totalSum - currentSum)
        {
            break;
        }
    }
    // 结果升序排列匹配输出要求
    result.Sort();
    return result;
}

逻辑验证 #

用题目给出的示例数组[3,7,5,6,2]测试：

• 数组总和为23，需要选中元素和>11.5即至少12
• 降序排序后数组为[7,6,5,3,2]
• 取第一个元素7，当前和为7，剩余和为16，不满足要求
• 取第二个元素6，当前和为13，剩余和为10，满足要求，停止选取
• 结果排序后为[6,7]，和示例正确输出完全一致

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