【Leetcode】805. Split Array With Same Average

题目地址：

https://leetcode.com/problems/split-array-with-same-average/description/

给定一个长$n$数组$A$，$A[i]\ge 0$。问是否能将其划分为两个非空子集，使得两个子集的平均数相等。$n\le 30,A[i]\le 10^4$。

设划分的两个子集分别为$B,C$，则有$$\frac{\sum B}{|B|}=\frac{\sum C}{|C|}=\frac{\sum B+\sum C}{|B|+|C|}=\frac{\sum A}{|A|}$$而$\frac{\sum A}{|A|}$与划分无关，所以我们其实就是要求是否存在一个子集，其平均值等于$\frac{\sum A}{|A|}$。考虑新数组$B[i]=A[i]-\frac{\sum A}{|A|}$，那么问题即转化为问$B$数组是否存在和为$0$的非平凡子集（非平凡的意思是不能为空，也不能为$B$自己）。为了避免处理浮点数，我们可以将$B$乘以一个系数，使其每个数都为整数。
这个是背包问题，通常来讲，这种题目可以用动态规划做，但此题的数据范围较为特殊，$n$较小，而$A[i]$的范围较大，如果用动态规划容易爆空间。可以考虑用双向DFS来做，参考https://blog.csdn.net/qq_46105170/article/details/115587834。具体思路是，我们先考虑前$n/2$个数，进行DFS，用哈希表存一下其所有子集的和以及每个子集的元素个数（这是为了排除掉平凡子集），如果前$n/2$个数就已经有非平凡子集和为$0$了，那就说明有解，可以提前结束；否则再暴力DFS后$n/2$个数的每个子集的和$s$，看哈希表是否存在$-s$并且元素个数总和大于$0$小于$n$，如果是，则有解。暴力枚举完毕没发现解，则无解。代码如下：

class Solution {
 public:
  bool splitArraySameAverage(vector<int> &A) {
    int sum = 0;
    for (int x : A) sum += x;
    int g = gcd(sum, A.size());
    int a = sum / g, b = A.size() / g;
    for (int &x : A) x = x * b - a;
    unordered_map<int, int> mp;
    return dfs1(0, 0, 0, A, mp) || dfs2(A.size() / 2, 0, 0, A, mp);
  }

  bool dfs2(int u, int cnt, int sum, vector<int> &A,
            unordered_map<int, int> &mp) {
    auto it = mp.find(-sum);
    if (it != mp.end() && it->second + cnt && it->second + cnt < A.size())
      return true;
    for (int i = u; i < A.size(); i++)
      if (dfs2(i + 1, cnt + 1, sum + A[i], A, mp)) return true;

    return false;
  }

  // 统计前一半的数的每个子集的和以及元素个数
  bool dfs1(int u, int cnt, int sum, vector<int> &A,
            unordered_map<int, int> &mp) {
    if (!sum && cnt && cnt < A.size()) return true;
    auto it = mp.find(sum);
    if (it == mp.end() || it->second > cnt) mp[sum] = cnt;
    for (int i = u; i < A.size() / 2; i++)
      if (dfs1(i + 1, cnt + 1, sum + A[i], A, mp)) return true;

    return false;
  }

  int gcd(int a, int b) { return b ? gcd(b, a % b) : a; }
};
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

时间复杂度$O(2^{n/2+1})$，空间$O(2^{n/2})$。