cswork / libraryOpen practice workspace

Coding questions · Amazon

Implement binary search lower/upper bounds

Implement lower_bound and upper_bound on a non-decreasing sorted array using iterative binary search in O(log n) time and O(1) space. This Amazon ML engineer screen tests precise boundary handling, loop-invariant reasoning, and edge cases like duplicates, empty arrays, and out-of-range targets.

Question

Given a non-decreasing sorted integer array nums of length n and a target value, implement two functions using binary search:

  1. lower_bound(nums, target) — return the smallest index i such that nums[i] >= target. Return n (i.e. len(nums)) if no such index exists.
  2. upper_bound(nums, target) — return the smallest index i such that nums[i] > target. Return n (i.e. len(nums)) if no such index exists.

Requirements:

  • Provide iterative implementations.
  • Achieve O(log n) time and O(1) extra space.
  • Handle the following edge cases: empty array; all elements strictly less than target; all elements strictly greater than target; duplicate values; negative numbers; and targets that fall outside the array’s value range.
  • Include a few test cases that demonstrate correctness.
  • Briefly argue correctness, explaining your boundary conditions and loop invariants.

Note that upper_bound(nums, target) - lower_bound(nums, target) gives the count of elements equal to target, and both functions return a valid insertion point that keeps the array sorted.

Binary Search: lower_bound

Given a non-decreasing sorted integer array nums of length n and a target value, implement lower_bound(nums, target) iteratively using binary search: return the smallest index i such that nums[i] >= target. If no such index exists, return n (i.e. len(nums)), which is the valid insertion point that keeps the array sorted.

Requirements:

  • Iterative implementation (no recursion).
  • O(log n) time, O(1) extra space.
  • Correctly handle: empty array; all elements strictly less than target; all elements strictly greater than target; duplicate values; negative numbers; and targets that fall outside the array's value range.

With duplicates of target present, lower_bound returns the index of the FIRST occurrence.

Constraints

  • 0 <= len(nums) <= 10^5
  • nums is sorted in non-decreasing order
  • -10^9 <= nums[i], target <= 10^9
  • Must be iterative with O(1) extra space

Examples

Example 1

Input: ([1, 2, 2, 2, 4, 5], 2)
Output: 1

Explanation: First index whose value is >= 2 is index 1 (the first of the three 2s).

Example 2

Input: ([1, 2, 2, 2, 4, 5], 3)
Output: 4

Explanation: 3 is absent; first value >= 3 is the 4 at index 4 — a valid insertion point.

Hints

  1. Use a half-open search interval [lo, hi) with lo = 0 and hi = len(nums). Initializing hi to n (not n-1) is what lets the answer legitimately equal n (the not-found / insert-at-end sentinel).
  2. Loop while lo < hi. When nums[mid] < target, the answer must be strictly right, so set lo = mid + 1; otherwise nums[mid] >= target so mid is a candidate, set hi = mid (do not exclude mid).
  3. Invariant: every index < lo fails the predicate (nums[i] < target) and every index >= hi satisfies it; the interval shrinks every step, so on exit lo == hi is the first index with nums[i] >= target.

Binary Search: upper_bound

Given a non-decreasing sorted integer array nums of length n and a target value, implement upper_bound(nums, target) iteratively using binary search: return the smallest index i such that nums[i] > target (strictly greater). If no such index exists, return n (i.e. len(nums)), which is the valid insertion point that keeps the array sorted.

Requirements:

  • Iterative implementation (no recursion).
  • O(log n) time, O(1) extra space.
  • Correctly handle: empty array; all elements strictly less than target; all elements strictly greater than target; duplicate values; negative numbers; and targets that fall outside the array's value range.

With duplicates of target present, upper_bound returns the index ONE PAST the last occurrence. Note upper_bound(nums, target) - lower_bound(nums, target) equals the count of elements equal to target. This implementation differs from lower_bound by exactly one comparison: <= instead of <.

Constraints

  • 0 <= len(nums) <= 10^5
  • nums is sorted in non-decreasing order
  • -10^9 <= nums[i], target <= 10^9
  • Must be iterative with O(1) extra space

Examples

Example 1

Input: ([1, 2, 2, 2, 4, 5], 2)
Output: 4

Explanation: Three 2s at indices 1-3; first value strictly > 2 is the 4 at index 4 (one past the last 2).

Example 2

Input: ([1, 2, 2, 2, 4, 5], 3)
Output: 4

Explanation: 3 absent; first value > 3 is the 4 at index 4. Same as lower_bound(3) here since 3 is missing.

Hints

  1. Reuse the lower_bound skeleton with a half-open interval [lo, hi), lo = 0, hi = len(nums). The only change needed is the predicate.
  2. Partition on the stricter predicate nums[i] > target. Elements equal to target must NOT count, so when nums[mid] <= target push lo = mid + 1; only when nums[mid] > target set hi = mid.
  3. Because <= sends equal elements to the right, upper_bound lands one index past the last occurrence of target, while lower_bound lands on the first — their difference is the count of target.

Solution

Both functions are the same half-open [lo, hi) binary search with hi initialized to len(nums) so the not-found/insertion sentinel n is reachable. The only difference is the comparison: lower_bound uses nums[mid] < target and upper_bound uses nums[mid] <= target, which makes upper_bound skip past elements equal to the target. The maintained invariant (everything left of lo fails the predicate, everything at/right of hi satisfies it) proves correctness, and halving the interval each step gives O(log n) time / O(1) space.

Both functions are variations of the same half-open-interval binary search. We maintain a search range [lo, hi) that is guaranteed to contain the answer, where lo starts at 0 and hi starts at n (one past the last index). The answer index can legitimately equal n (meaning “insert at the end / not found”), which is exactly why hi is initialized to n rather than n - 1.

def lower_bound(nums, target):
    # Smallest index i with nums[i] >= target, or len(nums).
    lo, hi = 0, len(nums)          # half-open [lo, hi)
    while lo < hi:
        mid = lo + (hi - lo) // 2  # avoids overflow in fixed-width langs
        if nums[mid] < target:
            lo = mid + 1           # nums[mid] too small; answer is to the right
        else:
            hi = mid               # nums[mid] >= target; answer is mid or left
    return lo

def upper_bound(nums, target):
    # Smallest index i with nums[i] > target, or len(nums).
    lo, hi = 0, len(nums)
    while lo < hi:
        mid = lo + (hi - lo) // 2
        if nums[mid] <= target:    # only difference: <= instead of <
            lo = mid + 1           # nums[mid] not strictly greater; go right
        else:
            hi = mid               # nums[mid] > target; answer is mid or left
    return lo

Why they differ by one comparison. lower_bound partitions on the predicate nums[i] >= target and pushes lo past everything that fails it (nums[mid] < target). upper_bound partitions on the stricter predicate nums[i] > target, so it must also push past elements that equal the target (nums[mid] <= target). Changing < to <= is the only change needed.

Correctness / loop invariant. Invariant: every index < lo satisfies the “too small” predicate, and every index >= hi satisfies the target predicate; the answer is the first index in [lo, hi) satisfying the predicate. Each iteration strictly shrinks hi - lo (either lo increases to mid + 1 > lo, or hi decreases to mid < hi), so the loop terminates. On exit lo == hi, and by the invariant lo is the smallest index satisfying the predicate — or n if none do.

Complexity. Each iteration halves the interval, giving O(log n) time. Only a constant number of integer variables are used, giving O(1) extra space. Both run iteratively with no recursion stack.

Edge cases (covered by the invariant, no special-casing needed):

  • Empty array (n == 0): lo == hi == 0, loop never runs, both return 0 == len(nums).
  • All elements < target: lo advances to n; both return n (not found).
  • All elements > target: both return 0.
  • Duplicates: lower_bound returns the index of the first occurrence; upper_bound returns one past the last occurrence; their difference is the count of target.
  • Negatives / out-of-range targets: handled identically — only </<= comparisons are used, so sign and magnitude are irrelevant.

Test cases:

nums = [1, 2, 2, 2, 4, 5]
assert lower_bound(nums, 2) == 1   # first index of 2
assert upper_bound(nums, 2) == 4   # one past last 2
assert upper_bound(nums, 2) - lower_bound(nums, 2) == 3  # count of 2s
assert lower_bound(nums, 3) == 4   # 3 absent; first index >= 3 is the 4
assert upper_bound(nums, 3) == 4
assert lower_bound(nums, 0) == 0   # target below range
assert lower_bound(nums, 9) == 6   # target above range -> n
assert upper_bound(nums, 9) == 6
assert lower_bound([], 5) == 0     # empty array
assert upper_bound([], 5) == 0

neg = [-5, -3, -3, 0, 4]
assert lower_bound(neg, -3) == 1
assert upper_bound(neg, -3) == 3
assert lower_bound(neg, -10) == 0
assert lower_bound(neg, 10) == 5