Return to Revisions

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