LeetCode 题解工作台

数组的均值分割

给定你一个整数数组 nums 我们要将 nums 数组中的每个元素移动到 A 数组 或者 B 数组中,使得 A 数组和 B 数组不为空,并且 average(A) == average(B) 。 如果可以完成则返回 true , 否则返回 false 。 注意: 对于数组 arr , average…

category

5

题型

code_blocks

4

代码语言

hub

3

相关题

当前训练重点

困难 · 状态·转移·动态规划

bolt

答案摘要

根据题目要求,要判断是否可以将数组 划分为两个子数组 和 ,使得两个子数组的平均值相等。 我们记数组 的和为 ,元素个数为 。子数组 的和以及个数分别为 和 ,那么子数组 的和为 $s_2 = s - s_1$,个数为 $n - k$,即:

Interview AiBox logo

Interview AiBox 实时 AI 助手,陪你讲清 状态·转移·动态规划 题型思路

试试 AiBox 面试助手arrow_forward
description

题目描述

给定你一个整数数组 nums

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

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

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

 

示例 1:

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

示例 2:

输入: nums = [3,1]
输出: false

 

提示:

  • 1 <= nums.length <= 30
  • 0 <= nums[i] <= 104
lightbulb

解题思路

方法一:折半查找 + 二进制枚举

根据题目要求,要判断是否可以将数组 nums\textit{nums} 划分为两个子数组 AABB,使得两个子数组的平均值相等。

我们记数组 nums\textit{nums} 的和为 ss,元素个数为 nn。子数组 AA 的和以及个数分别为 s1s_1kk,那么子数组 BB 的和为 s2=ss1s_2 = s - s_1,个数为 nkn - k,即:

s1k=s2nk=ss1nk\frac{s_1}{k} = \frac{s_2}{n - k} = \frac{s-s_1}{n-k}

整理可得:

s1×(nk)=(ss1)×ks_1 \times (n-k) = (s-s_1) \times k

化简可得:

s1k=sn\frac{s_1}{k} = \frac{s}{n}

也就是说,要我们找出一个子数组 AA,使得其平均值等于数组 nums\textit{nums} 的平均值。我们考虑将数组 nums\textit{nums} 每个元素都减去数组 nums\textit{nums} 的平均值,这样问题就转化为了在数组 nums\textit{nums} 中找出一个子数组,使得其和为 00

但是,数组 nums\textit{nums} 的平均值可能不是整数,浮点数计算可能存在精度问题,我们不妨将数组 nums\textit{nums} 中的每个元素都乘以 nn,即 nums[i]nums[i]×nnums[i] \leftarrow nums[i] \times n,上述式子就变成:

s1×nk=s\frac{s_1\times n}{k} = s

此时我们将数组 nums\textit{nums} 中每个元素都减去整数 ss,题目就转化为:在数组 numsnums 中找出一个子数组 AA,使得其和为 00

数组 nums\textit{nums} 的长度范围为 [1,30][1, 30],如果我们使用暴力枚举子数组的方法,时间复杂度为 O(2n)O(2^n),会超时。我们可以使用折半查找的方法,将时间复杂度降低到 O(2n/2)O(2^{n/2})

我们将数组 nums\textit{nums} 分成左右两部分,那么子数组 AA 可能存在三种情况:

  1. 子数组 AA 完全在数组 nums\textit{nums} 的左半部分;
  2. 子数组 AA 完全在数组 nums\textit{nums} 的右半部分;
  3. 子数组 AA 一部分在数组 nums\textit{nums} 的左半部分,一部分在数组 nums\textit{nums} 的右半部分。

我们可以使用二进制枚举的方法,先枚举左半部分所有子数组的和,如果存在一个子数组和为 00,直接返回 true,否则我们将其存入哈希表 vis\textit{vis} 中;然后枚举右半部分所有子数组的和,如果存在一个子数组和为 00,直接返回 true,否则我们判断此时哈希表 vis\textit{vis} 中是否存在该和的相反数,如果存在,直接返回 true

需要注意的是,我们不能同时全选左半部分和右半部分,因为这样会导致子数组 BB 为空,这是不符合题意的。在实现上,我们只需要考虑数组的 n1n-1 个数。

时间复杂度 O(n×2n2)O(n\times 2^{\frac{n}{2}}),空间复杂度 O(2n2)O(2^{\frac{n}{2}})

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
class Solution:
    def splitArraySameAverage(self, nums: List[int]) -> bool:
        n = len(nums)
        if n == 1:
            return False
        s = sum(nums)
        for i, v in enumerate(nums):
            nums[i] = v * n - s
        m = n >> 1
        vis = set()
        for i in range(1, 1 << m):
            t = sum(v for j, v in enumerate(nums[:m]) if i >> j & 1)
            if t == 0:
                return True
            vis.add(t)
        for i in range(1, 1 << (n - m)):
            t = sum(v for j, v in enumerate(nums[m:]) if i >> j & 1)
            if t == 0 or (i != (1 << (n - m)) - 1 and -t in vis):
                return True
        return False
speed

复杂度分析

指标
时间complexity is O(n * 2^n) in naive recursion but can be reduced using DP and bitmask optimization; space complexity is O(n * sum) to store achievable subset sums for each subset size.
空间Depends on the final approach
psychology

面试官常问的追问

外企场景
  • question_mark

    Looking for DP with careful state definition and sum tracking

  • question_mark

    Expect candidates to notice integer divisibility constraints

  • question_mark

    Interested in optimization via subset sum pruning or bitmask representation

warning

常见陷阱

外企场景
  • error

    Overlooking non-integer average feasibility for subset sizes

  • error

    Using naive recursion without pruning leading to TLE

  • error

    Failing to consider empty subsets, which are invalid

swap_horiz

进阶变体

外企场景
  • arrow_right_alt

    Find all possible splits rather than returning true/false

  • arrow_right_alt

    Allow splitting into more than two subarrays with the same average

  • arrow_right_alt

    Restrict subsets to contiguous elements while maintaining equal averages

help

常见问题

外企场景

数组的均值分割题解:状态·转移·动态规划 | LeetCode #805 困难