Skip to content

3134. Find the Median of the Uniqueness Array 👍

  • Time: $O(n\log n)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
class Solution {
 public:
  int medianOfUniquenessArray(vector<int>& nums) {
    const int n = nums.size();
    const long subarryCount = n * (n + 1L) / 2;
    const long medianCount = (subarryCount + 1) / 2;
    int l = 1;
    int r = n;

    while (l < r) {
      const int m = (l + r) / 2;
      if (subarrayWithAtMostKDistinct(nums, m) >= medianCount)
        r = m;
      else
        l = m + 1;
    }

    return l;
  }

 private:
  // Similar to 992. Subarrays with K Different Integers
  long subarrayWithAtMostKDistinct(const vector<int>& nums, int k) {
    long res = 0;
    unordered_map<int, int> count;

    for (int l = 0, r = 0; r < nums.size(); ++r) {
      if (++count[nums[r]] == 1)
        --k;
      while (k == -1)
        if (--count[nums[l++]] == 0)
          ++k;
      res += r - l + 1;  // nums[l..r], nums[l + 1..r], ..., nums[r]
    }

    return res;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
class Solution {
  public int medianOfUniquenessArray(int[] nums) {
    final int n = nums.length;
    final long subarrayCount = n * (n + 1L) / 2;
    final long medianCount = (subarrayCount + 1) / 2;
    int l = 1;
    int r = n;

    while (l < r) {
      final int m = (l + r) / 2;
      if (subarrayWithAtMostKDistinct(nums, m) >= medianCount)
        r = m;
      else
        l = m + 1;
    }

    return l;
  }

  // Similar to 992. Subarrays with K Different Integers
  private long subarrayWithAtMostKDistinct(int[] nums, int k) {
    long res = 0;
    HashMap<Integer, Integer> count = new HashMap<>();

    for (int l = 0, r = 0; r < nums.length; ++r) {
      if (count.merge(nums[r], 1, Integer::sum) == 1)
        --k;
      while (k == -1)
        if (count.merge(nums[l++], -1, Integer::sum) == 0)
          ++k;
      res += r - l + 1; // nums[l..r], nums[l + 1..r], ..., nums[r]
    }

    return res;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
class Solution:
  def medianOfUniquenessArray(self, nums: list[int]):
    n = len(nums)
    subarrayCount = n * (n + 1) // 2
    medianCount = (subarrayCount + 1) // 2

    # Similar to 992. Subarrays with K Different Integers
    def subarraysWithAtMostKDistinct(k: int) -> int:
      res = 0
      count = collections.Counter()

      l = 0
      for r, num in enumerate(nums):
        count[num] += 1
        if count[num] == 1:
          k -= 1
        while k < 0:
          count[nums[l]] -= 1
          if count[nums[l]] == 0:
            k += 1
          l += 1
        res += r - l + 1  # nums[l..r], nums[l + 1..r], ..., nums[r]

      return res

    return bisect.bisect_left(
        range(1, n), medianCount,
        key=lambda m: subarraysWithAtMostKDistinct(m)) + 1