Leetcode 805. 数组的均值分割 C++

Leetcode 805. 数组的均值分割

题目

给定的整数数组 A ，我们要将 A数组 中的每个元素移动到 B数组 或者 C数组中。（B数组和C数组在开始的时候都为空）

返回true ，当且仅当在我们的完成这样的移动后，可使得B数组的平均值和C数组的平均值相等，并且B数组和C数组都不为空。

示例:

输入: 
[1,2,3,4,5,6,7,8]
输出: true
解释: 我们可以将数组分割为 [1,4,5,8] 和 [2,3,6,7], 他们的平均值都是4.5。

注意:

• A 数组的长度范围为 [1, 30].
• A[i] 的数据范围为 [0, 10000].

题解

目标是B的均值=C的均值=A的均值，若A元素个数为n，B元素个数为k，那么就有(b1+b2+···+bk)/k = (a1+a2+···+an)/n，我们用乘法表示，即(b1+b2+···+bk)*n = sum(A)k，可得 b1 * n-sum(A) + b2 * n-sum(A) + ··· + bn * n-sum(A)=0。
因此，我们对A中元素进行处理，A[i] = nA[i]-sum(A)，我们只需要找到A的一个真子集B，其和为0。
我们不妨利用归并思想，分别对A的前半部分和后半部分看是否存在和为0的子集，有则直接输出true；然后再判断前部分的和以及后部分的和是否存在相反数，有则输出true。
其中找出前、后半段和的方法，我们现在要找A[0~n/2]部分的和，对于A[0],其和有只有A[0]；我们再看第i-1个元素时，假设有j个和{s1、s2、···、sj}，那么考虑A[i]时，可能出现新的和s1+A[i]、s2+A[i]、···、sj+A[i]，因为A[i]>=0的，所以我们将j个和以及j个新和可以根据归并思想进行排序去重。这样我们就能得到前或后半段所有可能和。
详细过程见代码

代码

	vector<int> temp;
    bool getSum(int l,int r,int sum,int& S,vector<int>& A,vector<int>& res){		//S存的可能和的个数，sum是一个目标和
        for(int i=l; i<r; i++){
            if(A[i]==0) return true;		//如果存在1个元素其为0，那么可以直接将其作为一个新集合，故返回true
            int j=0,k=0,n=0;
            while(j<S && k<S){		//归并思想索取新和
                if(res[j] < res[k]+A[i])    temp[n++] = res[j++];
                else if(res[j] == res[k]+A[i]){
                    if(res[j] == sum)   return true;
                    temp[n++] = res[j++];
                    k++;
                }else   temp[n++] = res[k++]+A[i];
            }
            while(j < S)    temp[n++] = res[j++];
            while(k < S)    temp[n++] = res[k++]+A[i];
            
            int m=0,m2=0;		//转存到答案res中，同时将只考虑单个元素A[i]加到集合中
            while(m2<n && temp[m2]<A[i])  res[m++] = temp[m2++];
            if(A[i] != temp[m2])    res[m++] = A[i];		
            else if(A[i] == sum)    return true;
            while(m2<n)  res[m++] = temp[m2++];
            S = m;
        }
        return false;
    }
    bool splitArraySameAverage(vector<int>& A) {
        int n = A.size(),sum=0;
        if(n <= 1)  return false;
        for(int i=0; i<n; i++)
            sum += A[i];
        for(int i=0; i<n; i++)
            A[i] = n*A[i]-sum;
        sort(A.begin(),A.end());
        vector<int> first((1<<(n/2))-1);
        vector<int> second((1<<(n/2+n%2))-1);
        temp = vector<int>(1<<(n/2+n%2));
        int sum1=0,fNum=0,sNum=0;
        for(int i=0; i<n/2; i++)    sum1 += A[i];
        if(getSum(0,n/2,sum1,fNum,A,first)) return true;		//获取前半段所有和的可能，只是如果可以直接得到一个和为0的集合，便可以直接输出
        if(getSum(n/2,n,-sum1,sNum,A,second))   return true;	//获取后半段所有和的可能，只是如果可以直接得到一个和为0的集合，便可以直接输出
        int l=0,r=sNum-1;
        while(l<fNum && r>=0){
            if(first[l]+second[r] > 0)  r--;
            else if(first[l]+second[r] < 0) l++;
            else{
                if(first[l]!=sum1 || second[r]!=-sum1)    return true;
                l++;
            }
        }
        return false;
    }

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/split-array-with-same-average
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。