Here is a cleaner and better way to solve the problem.
# Return the smallest index where the element is bigger than A[cur].
# if len(A) is returned, no element is bigger than A[cur].
def next_bigger_element(A, cur):
lo, hi = cur, len(A)
while lo + 1 < hi:
mid = (lo + hi)//2
if A[mid] == A[cur]:
lo = mid
else:
hi = mid
return hi
def distinct_elements_at_least(A, k):
cur = 0
count = 0
while cur != len(A) and count < k:
count += 1
cur = next_bigger_element(A, cur)
return count == k
To find whether A contains at least 4 distinct elements, just call distinct_elements_at_least(A, 4).
We can use this to solve all similar problems, such as whether A contains at least 2 distinct elements or whether A contains 777 distinct elements. The method works correctly even if you want to check whether A contains at least 0 element or whether A contains at least 1 element!
If you do not mind import bisect, you may like the following even shorter method, in the cases where, for example, elements in A are integers.
from bisect import bisect_left
def distinct_elements_at_least(A, k):
cur = 0
count = 0
while cur != len(A) and count < k:
count += 1
cur = bisect_left (A, A[cur] + 1, cur)
return count == k