Skip to content

2557. Maximum Number of Integers to Choose From a Range II

  • Time: $O(|\texttt{banned}|\log n)$
  • Space: $O(|\texttt{banned})$
 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 {
 public:
  int maxCount(vector<int>& banned, int n, long long maxSum) {
    const unordered_set<int> bannedSet{banned.begin(), banned.end()};
    int l = 1;
    int r = n;

    while (l < r) {
      const int m = (l + r + 1) / 2;
      if (getSum(bannedSet, m) > maxSum)
        r = m - 1;
      else
        l = m;
    }

    return l - ranges::count_if(bannedSet, [l](const int b) { return b <= l; });
  }

 private:
  // Returns sum([1..m]) - sum(bannedSet).
  long long getSum(const unordered_set<int>& bannedSet, long m) {
    long long sum = m * (m + 1) / 2;  // sum([1..m])
    for (const int b : bannedSet)
      if (b <= m)
        sum -= b;
    return sum;
  }
};

 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
class Solution {
  public int maxCount(int[] banned, int n, long maxSum) {
    Set<Integer> bannedSet = Arrays.stream(banned).boxed().collect(Collectors.toSet());
    int l = 1;
    int r = n;

    while (l < r) {
      final int m = (l + r + 1) / 2;
      if (getSum(bannedSet, m) > maxSum)
        r = m - 1;
      else
        l = m;
    }

    int ans = l;
    for (final int b : bannedSet)
      if (b <= l)
        --ans;
    return ans;
  }

  // Returns sum([1..m]) - sum(bannedSet).
  private long getSum(Set<Integer> bannedSet, long m) {
    long sum = m * (m + 1) / 2; // sum([1..m])
    for (final int b : bannedSet)
      if (b <= m)
        sum -= b;
    return sum;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
class Solution:
  def maxCount(self, banned: List[int], n: int, maxSum: int) -> int:
    bannedSet = set(banned)
    l = 1
    r = n

    while l < r:
      m = (l + r + 1) // 2
      if self._getSum(bannedSet, m) > maxSum:
        r = m - 1
      else:
        l = m

    return l - sum(b <= l for b in banned)

  # Returns sum([1..m]) - sum(bannedSet).
  def _getSum(self, bannedSet: Set[int], m: int) -> int:
    return m * (m + 1) // 2 - sum(b for b in bannedSet if b <= m)