【阿里淘天】3月15日暑期实习机试-第一题-连续非空子数组

连续非空子数组

题面

思路

正向求解的话，需要枚举所有的子数组，复杂度会来到\(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 = 0
    last_0 = last_1 = last_2 = -1
    for i in range(n):
        if arr[i] == 0:
            last_0 = i
        elif arr[i] == 1:
            last_1 = i
        else:
            last_2 = i
        if 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 = 0
    start = 0
    for i in range(n + 1):
        if i == n or arr[i] == 2:
            if start <= i - 1:
                L = i - start
                total = L * (L + 1) // 2
                current_zero = current_one = 0
                total_zero = total_one = 0
                for j in range(start, i):
                    num = arr[j]
                    if num == 0:
                        current_zero += 1
                        current_one = 0
# 每次新出现一个0，他都可以带来current_zero个新的子数组
                        total_zero += current_zero
                    else:  # 只能是1
                        current_one += 1
                        current_zero = 0
# 1同理
                        total_one += current_one
                mex2_count += (total - total_zero - total_one)
            start = i + 1
    
    # 计算mex1_count：包含至少一个0，不含1，可能含2的子数组数目
    mex1_count = 0
    start = 0
    for i in range(n + 1):
        if i == n or arr[i] == 1:
            if start <= i - 1:
                L = i - start
                total = L * (L + 1) // 2
                current_2 = 0
                total_2 = 0
                for j in range(start, i):
                    num = arr[j]
                    if num == 2:
                        current_2 += 1
                    else:
                        if current_2 > 0:
                            total_2 += current_2 * (current_2 + 1) // 2
                            current_2 = 0
                if current_2 > 0:
                    total_2 += current_2 * (current_2 + 1) // 2
                mex1_count += (total - total_2)
            start = i + 1
    
    total_sum = mex1_count * 1 + mex2_count * 2 + mex3_count * 3
    return total_sum

# 示例测试
arr = [0, 1, 2]
print(calculate_mex_sum(arr))  # 输出6

这里的思想本质上就是滑动窗口，但是和传统的滑动窗口写法有很大不同，传统的滑动窗口一般是求出最短的满足条件的窗口，这里要求出所有的（说是最长的也对），所以在写法上要做出比较大的改变。