Determine whether a sorted array contain at least 4 distinct elements in O(log n) time

Question

On one of my previous courseworks, I was faced with the following problem, which I think is unrealistic when using a direct / straightforward approach that usually algorithms have by leveraging certain data structures like dictionaries, etc.

Design an O(log n) algorithm whose input is a sorted list A. The algorithm should return true if A contains at least 4 distinct elements. Otherwise the algorithm should return false.

My lecturer came up with the following solution:

def ThreeDiff(A):
  if A[0]==A[len(A)-1]:
    return -1

  minind=0
  maxind=len(A)-1

  while maxind-mind>1:
    midind=int((minind+maxind)/2)
    if A[midind]>A[minind] and A[midind]<A[maxind]:
      return midind
    if A[midind]==A[minind]:
      minind=midind
    else:
     maxind=midind
  return -1

def FourDiff(A):
  midind=ThreeDiff(A)
  if midind==-1:
    return false
  return ThreeDiff(A[0:midind+1])!=- 1 or ThreeDiff(A[midind:len(A)]) != -1

Is there a cleaner or better way to solve this problem?

Answer 1

Here is a cleaner and better way to solve the problem.

# Return the smallest index where the element is bigger than `A[start_index]`.
# If `len(A)` is returned, no element is bigger than `A[start_index]`.
def next_bigger_element(start_index, A):
    lo, hi = start_index, len(A)
    while lo + 1 < hi:
        mid = (lo + hi) // 2
        if A[mid] == A[start_index]:
            lo = mid
        else:
            hi = mid
    return hi


def distinct_elements_at_least(k, A):
    if len(A) == 0:
        return k <= 0
    index = 0
    count = 1
    # keep finding the next bigger element until `k` elements have
    # been found or we have reached the end of the array.
    while count < k and A[index] != A[-1]:
        index = next_bigger_element(index, A)
        count += 1
    return count >= k

To find whether A contains at least 4 distinct elements, just call distinct_elements_at_least(4, A).

This program works correctly for any given number k. For example, it can be used to check whether A has 0 element or whether A has 7 distinct elements. For any fixed k, it works in $O(\log n)$ time as at most k binary searches on an interval of size at most n are done.

---

If you do not mind import bisect, you may prefer the following shorter code, since method next_bigger_element is no longer needed.

from bisect import bisect_right

def distinct_elements_at_least(k, A):
    if len(A) == 0:
        return k <= 0
    index = 0
    count = 0
    while index < len(A) and count < k:
        count += 1
        index = bisect_right(A, A[index], index + 1)
    return count >= k