题目截图

题目分析

• 找到k的位置
• 然后一步步往左走，一步步往右走
• 统计左边和右边的比当前k小的和比k大的
• lst = [[small, big]]，分为left和right两部分
• 可以先一侧的单独看small和big，找到big - small = 0或者1的即可
• 如果两侧的话，比较麻烦
• 我们将left和right两侧的再转化成diff = big - small
• 那么我们需要在左侧找一个diff1和右侧找一个diff2
• 使得diff1 + diff2 = 0或者1即可
• 特别地，我们固定diff1，然后可以得到diff2两个可能的值
• 因此，可以先对diff2排序，再用bisect_left和bisect_right卡住所有可选的值即可

ac code

class Solution:
    def countSubarrays(self, nums: List[int], k: int) -> int:
        n = len(nums)
        k_idx = 0
        p = [0] * (n + 1)
        for i, v in enumerate(nums):
            p[v] = i
        
        while nums[k_idx] != k:
            k_idx += 1
        # s个比k小，s个比k大
        # s个比k小，s + 1个比k大
        left, right = k - 1, n - k
        lefts, rights = [], []
        # 左边的情况
        small, big = 0, 0
        for i in range(k_idx - 1, -1, -1):
            if nums[i] < k:
                small += 1
            else:
                big += 1
            lefts.append([small, big])
        # 右边的情况
        small, big = 0, 0
        for i in range(k_idx + 1, n):
            if nums[i] < k:
                small += 1
            else:
                big += 1
            rights.append([small, big])
        ans = 1 # 自己
        # small - big
        #print(lefts)
        #print(rights)
        lefts_diff = []
        for small, big in lefts:
            lefts_diff.append(big - small)
        rights_diff = []
        for small, big in rights:
            rights_diff.append(big - small)
        # print(lefts_diff)
        # print(rights_diff)
        # 两个和要等于0或者1
        ans = 1 # 自己
        for diff in lefts_diff:
            if diff == 0 or diff == 1:
                ans += 1
        for diff in rights_diff:
            if diff == 0 or diff == 1:
                ans += 1
        rights_diff.sort()
        for num in lefts_diff:
            need1 = -num
            need2 = 1 - num
            i = bisect_left(rights_diff, need1)
            j = bisect_right(rights_diff, need2)
            ans += j - i
        return ans

总结

• 一个关键点：比k大的 - 比k小的 = 0或者1
• 为了研究子数组，我们可以考虑记录左侧连续的和右侧连续的出现的small和big个数
• 我们需要big1 + big2 - small1 - small2 = 0或者1
• 直接先变成diff1和diff2使得，diff1 + diff2 = 0或者1
• 那就好办了
• 遍历diff1，排序diff2，用二分卡住两边，得到数量即可