3134. 找出唯一性数组的中位数

来源第 395 场周赛 Q4难度困难分数2451

题目描述

给你一个整数数组 nums 。数组 nums 的 唯一性数组 是一个按元素从小到大排序的数组，包含了 nums 的所有非空子数组中不同元素的个数。

换句话说，这是由所有 0 <= i <= j < nums.length 的 distinct(nums[i..j]) 组成的递增数组。

其中，distinct(nums[i..j]) 表示从下标 i 到下标 j 的子数组中不同元素的数量。

返回 nums 唯一性数组 的 中位数 。

注意，数组的 中位数 定义为有序数组的中间元素。如果有两个中间元素，则取值较小的那个。

示例 1：

输入：nums = [1,2,3]

输出：1

解释：

nums 的唯一性数组为 [distinct(nums[0..0]), distinct(nums[1..1]), distinct(nums[2..2]), distinct(nums[0..1]), distinct(nums[1..2]), distinct(nums[0..2])]，即 [1, 1, 1, 2, 2, 3] 。唯一性数组的中位数为 1 ，因此答案是 1 。

示例 2：

输入：nums = [3,4,3,4,5]

输出：2

解释：

nums 的唯一性数组为 [1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3] 。唯一性数组的中位数为 2 ，因此答案是 2 。

示例 3：

输入：nums = [4,3,5,4]

输出：2

解释：

nums 的唯一性数组为 [1, 1, 1, 1, 2, 2, 2, 3, 3, 3] 。唯一性数组的中位数为 2 ，因此答案是 2 。

提示：

1 <= nums.length <= 10⁵
1 <= nums[i] <= 10⁵

解法

方法一：二分查找 + 双指针

思考

唯一性数组含全部 \(O(n^2)\) 个子数组的不同元素个数，显式构造后再取中位数不可行。

子数组不同元素个数 \(\le x\) 的个数随 \(x\) 单调，中位数即第一个使该个数达到一半的 \(x\)。滑动窗口可在线性时间内统计「不同元素 \(\le mx\)」的子数组数。

因此对 \(x\) 二分，窗口右端纳入新值，当种类超过 \(mx\) 时左端右移，累加 \(r-l+1\)。个数达到 \(\lceil m/2\rceil\) 即合法。

我们记数组 \(\textit{nums}\) 的长度为 \(n\)，那么唯一性数组的长度为 \(m = \frac{(1 + n) \times n}{2}\)，而唯一性数组的中位数就是这 \(m\) 个数中的第 \(\frac{m + 1}{2}\) 小的数字。

考虑唯一性数组中，有多少个数小于等于 \(x\)。随着 \(x\) 的增大，只会有越来越多的数小于等于 \(x\)。这存在着单调性，因此，我们可以二分枚举 \(x\)，找到第一个 \(x\)，满足唯一性数组中小于等于 \(x\) 的数的个数大于等于 \(\frac{m + 1}{2}\)，这个 \(x\) 就是唯一性数组的中位数。

我们定义二分查找的左边界 \(l = 0\)，右边界 \(r = n\)，然后进行二分查找，对于每个 \(\textit{mid}\)，我们检查唯一性数组中小于等于 \(\textit{mid}\) 的数的个数是否大于等于 \(\frac{m + 1}{2}\)。我们通过函数 \(\text{check}(mx)\) 来实现这一点。

函数 \(\text{check}(mx)\) 的实现思路如下：

由于子数组越长，不同元素的个数越多，因此，我们可以利用双指针维护一个滑动窗口，使得窗口中的子数组的不同元素的个数不超过 \(mx\)。具体地，我们维护一个哈希表 \(\textit{cnt}\)，\(\textit{cnt}[x]\) 表示窗口中元素 \(x\) 的个数。我们使用两个指针 \(l\) 和 \(r\)，其中 \(l\) 表示窗口的左边界，而 \(r\) 表示窗口的右边界。初始时 \(l = r = 0\)。

我们枚举 \(r\)，对于每个 \(r\)，我们将 \(\textit{nums}[r]\) 加入窗口中，并更新 \(\textit{cnt}[\textit{nums}[r]]\)。如果窗口中的不同元素的个数超过了 \(mx\)，我们需要将 \(l\) 右移，直到窗口中的不同元素的个数不超过 \(mx\)。此时，右端点为 \(r\)，而左端点为 \([l,..r]\) 的子数组都是满足条件的，一共有 \(r - l + 1\) 个子数组。我们将这个数量累加到 \(k\) 中，如果 \(k\) 大于等于 \(\frac{m + 1}{2}\)，那么说明唯一性数组中小于等于 \(\textit{mid}\) 的数的个数大于等于 \(\frac{m + 1}{2}\)，我们返回 \(\text{true}\)，否则返回 \(\text{false}\)。

时间复杂度 \(O(n \times \log n)\)，空间复杂度 \(O(n)\)。其中 \(n\) 为数组 \(\textit{nums}\) 的长度。

Python3JavaC++GoTypeScript

class Solution:
    def medianOfUniquenessArray(self, nums: List[int]) -> int:
        def check(mx: int) -> bool:
            cnt = defaultdict(int)
            k = l = 0
            for r, x in enumerate(nums):
                cnt[x] += 1
                while len(cnt) > mx:
                    y = nums[l]
                    cnt[y] -= 1
                    if cnt[y] == 0:
                        cnt.pop(y)
                    l += 1
                k += r - l + 1
                if k >= (m + 1) // 2:
                    return True
            return False

        n = len(nums)
        m = (1 + n) * n // 2
        return bisect_left(range(n), True, key=check)

class Solution {
    private long m;
    private int[] nums;

    public int medianOfUniquenessArray(int[] nums) {
        int n = nums.length;
        this.nums = nums;
        m = (1L + n) * n / 2;
        int l = 0, r = n;
        while (l < r) {
            int mid = (l + r) >> 1;
            if (check(mid)) {
                r = mid;
            } else {
                l = mid + 1;
            }
        }
        return l;
    }

    private boolean check(int mx) {
        Map<Integer, Integer> cnt = new HashMap<>();
        long k = 0;
        for (int l = 0, r = 0; r < nums.length; ++r) {
            int x = nums[r];
            cnt.merge(x, 1, Integer::sum);
            while (cnt.size() > mx) {
                int y = nums[l++];
                if (cnt.merge(y, -1, Integer::sum) == 0) {
                    cnt.remove(y);
                }
            }
            k += r - l + 1;
            if (k >= (m + 1) / 2) {
                return true;
            }
        }
        return false;
    }
}

class Solution {
public:
    int medianOfUniquenessArray(vector<int>& nums) {
        int n = nums.size();
        using ll = long long;
        ll m = (1LL + n) * n / 2;
        int l = 0, r = n;
        auto check = [&](int mx) -> bool {
            unordered_map<int, int> cnt;
            ll k = 0;
            for (int l = 0, r = 0; r < n; ++r) {
                int x = nums[r];
                ++cnt[x];
                while (cnt.size() > mx) {
                    int y = nums[l++];
                    if (--cnt[y] == 0) {
                        cnt.erase(y);
                    }
                }
                k += r - l + 1;
                if (k >= (m + 1) / 2) {
                    return true;
                }
            }
            return false;
        };
        while (l < r) {
            int mid = (l + r) >> 1;
            if (check(mid)) {
                r = mid;
            } else {
                l = mid + 1;
            }
        }
        return l;
    }
};

func medianOfUniquenessArray(nums []int) int {
    n := len(nums)
    m := (1 + n) * n / 2
    return sort.Search(n, func(mx int) bool {
        cnt := map[int]int{}
        l, k := 0, 0
        for r, x := range nums {
            cnt[x]++
            for len(cnt) > mx {
                y := nums[l]
                cnt[y]--
                if cnt[y] == 0 {
                    delete(cnt, y)
                }
                l++
            }
            k += r - l + 1
            if k >= (m+1)/2 {
                return true
            }
        }
        return false
    })
}

function medianOfUniquenessArray(nums: number[]): number {
    const n = nums.length;
    const m = Math.floor(((1 + n) * n) / 2);
    let [l, r] = [0, n];
    const check = (mx: number): boolean => {
        const cnt = new Map<number, number>();
        let [l, k] = [0, 0];
        for (let r = 0; r < n; ++r) {
            const x = nums[r];
            cnt.set(x, (cnt.get(x) || 0) + 1);
            while (cnt.size > mx) {
                const y = nums[l++];
                cnt.set(y, cnt.get(y)! - 1);
                if (cnt.get(y) === 0) {
                    cnt.delete(y);
                }
            }
            k += r - l + 1;
            if (k >= Math.floor((m + 1) / 2)) {
                return true;
            }
        }
        return false;
    };
    while (l < r) {
        const mid = (l + r) >> 1;
        if (check(mid)) {
            r = mid;
        } else {
            l = mid + 1;
        }
    }
    return l;
}

3134. 找出唯一性数组的中位数

题目描述

解法

方法一：二分查找 + 双指针

评论