10_flags.py

"""https://app.codility.com/programmers/lessons/10-prime_and_composite_numbers/flags/"""

from math import floor, sqrt

# Time complexity:
# If N is the length of array A then finding peaks = O(N),
# binary search is in the worst case O(log N) iterations,
# each iteration requiring a linear scan of the array in the worst case = O(N)
# Overall: O(N + N log N) = O(N log N)

def solution_a(a: list[int]) -> int:
    """
    Find maximum number of flags that can be set on peaks of array A.
    A peak is an index P such that 0 < P < N-1 and A[P-1] < A[P] > A[P+1].
    Flags can only be set on peaks, and the distance between any two flags
    must be at least the number of flags set.
    """

    # My solution that passes 100% of tests but is not optimal

    n = len(a)

    # Find peaks
    peaks = []
    for i in range(1, n-1):
        if a[i-1] < a[i] > a[i+1]:
            peaks.append(i)

    def can_place_flags(f):
        """Greedily set flags whenever possible for maximum flexibility in later placement"""
        if not peaks:
            return False
        placed = 1
        last_placed = peaks[0]
        for i in range(1, len(peaks)):
            if peaks[i] - last_placed >= f:
                last_placed = peaks[i]
                placed += 1
        return placed >= f

    # Binary search on the range (no points are peaks, all points are peaks)
    l, r = 0, len(peaks)
    res = 0
    while l <= r:
        m = (l+r) // 2
        if can_place_flags(m):
            res = m
            l = m + 1
        else:
            r = m - 1
    return res

# Time complexity:
# Precompute next_peak array = O(N)
# Our outer loop only checks up to sqrt(N) flags since we cannot place more,
# given the distance requirement.
# The inner check is bounded by its input size, since it exits after placing f flags,
# which we assert is at most sqrt(N).
# Overall: O(N + sqrt(N) * sqrt(N)) = O(N + N) = O(N)

def solution_b(a: list[int]) -> int:
    """Optimised solution generated by gemini-3-pro (explanations added)"""

    n = len(a)

    if n == 0:
        return 0

    # Precompute next_peak array
    next_peak = [-1] * n
    next_p = -1
    for i in range(n-2, -1, -1):
        if i > 0 and a[i-1] < a[i] > a[i+1]:
            next_p = i
        next_peak[i] = next_p

    def can_place_flags(f):
        """Use next_peak array to jump to next valid peak position"""
        pos = 0
        placed = 0
        while pos < n and placed < f: # Worst case f iterations (since f < N peaks)
            pos = next_peak[pos]
            if pos == -1:
                break
            placed += 1
            pos += f
        return placed >= f

    # Only check up to sqrt(N) flags as that's the maximum possible, given that
    # min distance between flags is the number of flags.

    i = 1
    max_flags = 0
    # Since dist between flags is at least the number of flags,
    # it might be possible to place 1 more than sqrt(N) flags,
    while i <= floor(sqrt(n)) + 1:
        if can_place_flags(i):
            max_flags = i
        i += 1
    return max_flags