Binary Search: First & Last Occurrences and Element Counting

This chapter explores how to leverage Binary Search to find the first and last occurrences of a target element in a sorted array, and how to extend this logic to count the total occurrences of an element in $O(\log N)$ time. We will transition from a naive linear scan to an optimized approach using Lower/Upper Bounds, and finally build a custom, raw Binary Search solution from scratch.

1. The Problem Statement & Core Intuition

Given a sorted array of integers containing duplicate values and a target value $X$, we need to solve two closely related problems:

1. First and Last Occurrences: Find the starting (first) and ending (last) index of $X$. If $X$ is not present in the array, return [-1, -1].
2. Count Occurrences: Find the total number of times $X$ appears in the array.

Real-World Example

Consider the sorted array:

2. The Naive Approach (Linear Scan)

The most straightforward way to find the first and last occurrences of $X$ is to traverse the array from left to right.

Algorithm Workflow

1. Initialize first = -1 and last = -1.
2. Iterate through the array from index $0$ to $N-1$:
   • If the current element arr[i] == X:
     • If first is still -1, this is the first time we are seeing the element. Update first = i.
     • Always update last = i because any subsequent match will push the last occurrence further to the right.
3. Return the pair [first, last].

Code Implementation (C++)

#include <vector>
#include <utility>

std::pair<int, int> findFirstAndLastNaive(const std::vector<int>& arr, int n, int x) {
    int first = -1;
    int last = -1;
    for (int i = 0; i < n; i++) {
        if (arr[i] == x) {
            if (first == -1) {
                first = i;
            }
            last = i;
        }
    }
    return {first, last};
}

Complexity Analysis

• Time Complexity: $\mathcal{O}(N)$ because we must scan the entire array of size $N$.
• Space Complexity: $\mathcal{O}(1)$ as we only use a few auxiliary variables.

3. Mental Model: Lower Bound & Upper Bound Connection

Since the input array is sorted, we can optimize our search from $\mathcal{O}(N)$ to $\mathcal{O}(\log N)$ using Binary Search. We can achieve this by reusing two fundamental Binary Search concepts: Lower Bound and Upper Bound.

Definitions

• Lower Bound (LB): The smallest index i such that $\text{arr}[i] \ge X$.
• Upper Bound (UB): The smallest index i such that $\text{arr}[i] > X$.

The Mathematical Relationship

If the target element $X$ exists in the array:

• The first occurrence of $X$ will always be at the index returned by the Lower Bound of $X$.
• The last occurrence of $X$ will always be at the index immediately preceding the Upper Bound of $X$ (i.e., $\text{UB} - 1$).

4. Optimized Approach 1: Lower & Upper Bound Implementation

To implement this approach, we first define our helper functions for Lower Bound and Upper Bound, and then handle the edge cases where the element $X$ does not exist in the array.

Critical Edge Cases

1. Element is larger than any element in the array: The Lower Bound will point to the hypothetical index $N$. We cannot access $\text{arr}[N]$.
2. Element is absent but falls within the array's range: The Lower Bound will point to the first element greater than $X$. For example, if we search for $10$ in the array above, the Lower Bound points to index $6$ (value $11$). Since $\text{arr}[6] \neq 10$, the element is absent.

Thus, the element is present if and only if:

Code Implementation (C++)

#include <vector>
#include <utility>

// Returns the smallest index where arr[index] >= x
int getLowerBound(const std::vector<int>& arr, int n, int x) {
    int low = 0, high = n - 1;
    int ans = n;
    while (low <= high) {
        int mid = low + (high - low) / 2;
        if (arr[mid] >= x) {
            ans = mid;
            high = mid - 1; // Look for smaller index on the left
        } else {
            low = mid + 1;  // Look on the right
        }
    }
    return ans;
}

// Returns the smallest index where arr[index] > x
int getUpperBound(const std::vector<int>& arr, int n, int x) {
    int low = 0, high = n - 1;
    int ans = n;
    while (low <= high) {
        int mid = low + (high - low) / 2;
        if (arr[mid] > x) {
            ans = mid;
            high = mid - 1; // Look for smaller index on the left
        } else {
            low = mid + 1;  // Look on the right
        }
    }
    return ans;
}

std::pair<int, int> firstAndLastOccurrence(const std::vector<int>& arr, int n, int x) {
    int lb = getLowerBound(arr, n, x);
    
    // Check if element is not present
    if (lb == n || arr[lb] != x) {
        return {-1, -1};
    }
    
    int ub = getUpperBound(arr, n, x);
    return {lb, ub - 1};
}

Complexity Analysis

• Time Complexity: $\mathcal{O}(2 \log_2 N) \approx \mathcal{O}(\log_2 N)$ because we perform two independent binary searches.
• Space Complexity: $\mathcal{O}(1)$ as no extra memory is allocated.

5. Optimized Approach 2: Raw Binary Search from Scratch

Some interviewers may request that you write the binary search logic directly without relying on Lower Bound or Upper Bound abstractions. We can write two distinct, clean binary search routines: one for the first occurrence and one for the last occurrence.

1. First Occurrence Logic

• If arr[mid] == X:
  • This is a potential candidate for the first occurrence. Record first = mid.
  • Since we want the first occurrence, we must continue searching the left half to see if $X$ appears earlier. Thus, we eliminate the right half: high = mid - 1.
• If arr[mid] < X: Search the right half (low = mid + 1).
• If arr[mid] > X: Search the left half (high = mid - 1).

2. Last Occurrence Logic

• If arr[mid] == X:
  • This is a potential candidate for the last occurrence. Record last = mid.
  • Since we want the last occurrence, we must continue searching the right half to see if $X$ appears later. Thus, we eliminate the left half: low = mid + 1.
• If arr[mid] < X: Search the right half (low = mid + 1).
• If arr[mid] > X: Search the left half (high = mid - 1).

Optimization Tip: If the first occurrence search returns -1, the element does not exist in the array. We can immediately return [-1, -1] without running the second binary search, saving unnecessary operations.

6. Visualizing the Binary Search Traversal

Let's trace the raw binary search for the first occurrence of $X = 8$ in the array:

First Occurrence Trace

7. Complete Raw Binary Search Code

Below is the clean, decoupled implementation of the raw binary search approach.

#include <vector>
#include <utility>

int findFirst(const std::vector<int>& arr, int n, int x) {
    int low = 0, high = n - 1;
    int first = -1;
    while (low <= high) {
        int mid = low + (high - low) / 2;
        if (arr[mid] == x) {
            first = mid;
            high = mid - 1; // Look for earlier occurrences on the left
        } else if (arr[mid] < x) {
            low = mid + 1;
        } else {
            high = mid - 1;
        }
    }
    return first;
}

int findLast(const std::vector<int>& arr, int n, int x) {
    int low = 0, high = n - 1;
    int last = -1;
    while (low <= high) {
        int mid = low + (high - low) / 2;
        if (arr[mid] == x) {
            last = mid;
            low = mid + 1; // Look for later occurrences on the right
        } else if (arr[mid] < x) {
            low = mid + 1;
        } else {
            high = mid - 1;
        }
    }
    return last;
}

std::pair<int, int> firstAndLastOccurrenceRaw(const std::vector<int>& arr, int n, int x) {
    int first = findFirst(arr, n, x);
    if (first == -1) {
        return {-1, -1}; // Optimization: If first doesn't exist, last won't either
    }
    int last = findLast(arr, n, x);
    return {first, last};
}

8. Problem 2: Counting Occurrences in a Sorted Array

Once we have solved the first and last occurrence problem, counting the total occurrences of $X$ becomes a trivial constant-time calculation.

The Formula

If the element $X$ exists in the array:

If the element does not exist (i.e., first == -1), the count is $0$.

Code Implementation (C++)

#include <vector>
#include <utility>

int countOccurrences(const std::vector<int>& arr, int n, int x) {
    std::pair<int, int> occurrences = firstAndLastOccurrenceRaw(arr, n, x);
    
    if (occurrences.first == -1) {
        return 0;
    }
    
    return occurrences.second - occurrences.first + 1;
}

9. Performance & Complexity Comparison

Summary and Key Takeaways

• Leverage Sorted Properties: Whenever you need to search or count elements in a sorted array, always think of Binary Search ($\mathcal{O}(\log N)$) instead of Linear Search ($\mathcal{O}(N)$).
• Lower Bound vs. Upper Bound:
  • The first occurrence of a target $X$ is always at $\text{Lower Bound}(X)$.
  • The last occurrence of a target $X$ is always at $\text{Upper Bound}(X) - 1$.
• Edge Case Validation: Always check if the Lower Bound index is valid (i.e., $\text{LB} < N$) and that the element at that index actually matches the target ($\text{arr}[\text{LB}] == X$) before assuming the element is present.
• Early Exit Optimization: When searching for first and last occurrences, if the first occurrence returns -1, skip the search for the last occurrence entirely to save $\mathcal{O}(\log N)$ operations.
• The Counting Formula: The total count of an element in a sorted array is calculated as $\text{last} - \text{first} + 1$. This reduces the counting problem to a simple coordinate subtraction after finding the boundaries.