连续非空子数组
题面
思路
正向求解的话,需要枚举所有的子数组,复杂度会来到\(O(n^3)\),完全不可行,在观察题目输入描述,\(a_i\)的取值范围非常小,故我们考虑反向求解(这也是非常经典的思路,无法直接计数,我们就计算每个答案的贡献是多少)
利用类似滑动窗口的思想,去统计\(mex(1), mex(2), mex(3)\)的贡献,\(mex(0)\)的贡献是0,无需统计。
代码
def calculate_mex_sum(arr):n = len(arr)# 计算mex3_count:包含0、1、2的子数组数目mex3_count = 0last_0 = last_1 = last_2 = -1for i in range(n):if arr[i] == 0:last_0 = ielif arr[i] == 1:last_1 = ielse:last_2 = iif last_0 != -1 and last_1 != -1 and last_2 != -1:s = min(last_0, last_1, last_2)mex3_count += s + 1# 计算mex2_count:包含0和1,不含2的子数组数目mex2_count = 0start = 0for i in range(n + 1):if i == n or arr[i] == 2:if start <= i - 1:L = i - starttotal = L * (L + 1) // 2current_zero = current_one = 0total_zero = total_one = 0for j in range(start, i):num = arr[j]if num == 0:current_zero += 1current_one = 0
# 每次新出现一个0,他都可以带来current_zero个新的子数组total_zero += current_zeroelse: # 只能是1current_one += 1current_zero = 0
# 1同理total_one += current_onemex2_count += (total - total_zero - total_one)start = i + 1# 计算mex1_count:包含至少一个0,不含1,可能含2的子数组数目mex1_count = 0start = 0for i in range(n + 1):if i == n or arr[i] == 1:if start <= i - 1:L = i - starttotal = L * (L + 1) // 2current_2 = 0total_2 = 0for j in range(start, i):num = arr[j]if num == 2:current_2 += 1else:if current_2 > 0:total_2 += current_2 * (current_2 + 1) // 2current_2 = 0if current_2 > 0:total_2 += current_2 * (current_2 + 1) // 2mex1_count += (total - total_2)start = i + 1total_sum = mex1_count * 1 + mex2_count * 2 + mex3_count * 3return total_sum# 示例测试
arr = [0, 1, 2]
print(calculate_mex_sum(arr)) # 输出6
这里的思想本质上就是滑动窗口,但是和传统的滑动窗口写法有很大不同,传统的滑动窗口一般是求出最短的满足条件的窗口,这里要求出所有的(说是最长的也对),所以在写法上要做出比较大的改变。