LC805. 数组的均值分割

Posted on | Post modified:

805. 数组的均值分割

题面

给定你一个整数数组 nums

我们要将 nums 数组中的每个元素移动到 A 数组 或者 B 数组中，使得 A 数组和 B 数组不为空，并且 average(A) == average(B) 。

如果可以完成则返回true ， 否则返回 false 。

注意：对于数组 arr , average(arr) 是 arr 的所有元素的和除以 arr 长度。

数据范围

  • $1 \leq n \leq 30$
  • $0 \leq nums[i] \leq 10^4$

题解

  • 由于$1 \leq n \leq 30$, 考虑折半搜索。

  • 分别对左侧数组和右侧数组进行二进制枚举，假设左侧选择了 $lc$ 个数字，和为 $ls$, 右侧选择了 $rc$ 个数字，和为$rs$，假设数组和为 $S$, 若要满足平均值相等，则有:

  • 化简可得：$S\times lc - n \times ls$ = $-(S \times rc - n \times rs)$

  • 预处理 $S\times lc - n \times ls$ 的值

  • 注意边界情况，左侧和右侧不能同时全选

代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
class Solution {
public:
    bool splitArraySameAverage(vector<int>& nums) {
        vector<int> L, R;
        int n = (int)nums.size();
        if(n <= 1) return false;
        int s = 0;
        for(int i = 0; i < n; i++) s += nums[i];
        if(n == 0) return true;
        if(n == 1) return false;
        for(int i = 0; i < n / 2; i++){
            L.push_back(nums[i]);
        }
        for(int i = n / 2; i < n; i++){
            R.push_back(nums[i]);
        }
        int ln = (int)L.size();
        int rn = (int)R.size();
        vector<pair<int, int>> A, B;
        for(int i = 0; i < (1 << ln); i++){
            int ss = 0;
            int cnt = 0;
            for(int j = 0; j < ln; j++){
                if((i >> j) & 1){
                    ss += L[j];
                    cnt++;
                }
            }
            A.push_back({ss, cnt});
        }

        for(int i = 0; i < (1 << rn); i++){
            int ss = 0;
            int cnt = 0;
            for(int j = 0; j < rn; j++){
                if((i >> j) & 1){
                    ss += R[j];
                    cnt++;
                }
            }
            B.push_back({ss, cnt});
        }
        map<int, int> rec, allL, noL;
        for(auto &v: A){
            int x1 = v.first;
            int y1 = v.second;
            if(y1 == 0){
                 noL[0] = 1;
                 continue;
            }
            if(y1 == ln){
                allL[s * y1 - n * x1] = 1;
                continue;
            }
            rec[s * y1 - n * x1] = 1;
            
        }
        for(auto &v: B){
            int x2 = v.first;
            int y2 = v.second;
            if(y2 != rn){
                if(allL[-(s * y2 - n * x2)] || rec[-(s * y2 - n * x2)]) return true;
            }
            else{
                if(noL[-(s * y2 - n * x2)] || rec[-(s * y2 - n * x2)]) return true;
            }
        }
        return false;
    }
};