#include <iostream>
#include <cstddef>
#include <cassert>

/*************************************************************
 * TWO SEARCH PROBLEMS
 *
 * Problem 1:
 *
 * Given a sorted array of distinct integers, find the
 * smallest index k such that
 *
 *     array[k] == k.
 *
 * If no such index exists, return capacity.
 *
 * Because the array is strictly increasing, we can use
 * a modified binary search.
 *
 *
 * Problem 2:
 *
 * Given an unsorted array, determine whether the entry
 * at a specified index would equal that index after
 * sorting the array into nondecreasing order.
 *
 * Entries need not be distinct.
 *
 * We can answer this question with one linear scan,
 * without sorting or modifying the array.
 *************************************************************/

// Function declarations

int main();

std::size_t equal_to_index(
    int const         array[],
    std::size_t const capacity
);

bool equal_to_index(
    std::size_t const idx,
    int const         array[],
    std::size_t const capacity
);

/*************************************************************
 * MAIN
 *************************************************************/

int main() {
    // Sorted arrays containing distinct entries.
    int sarray1[10]{
        -10, 0, 4, 5, 6, 9, 10, 12, 15, 19
    };

    int sarray2[10]{
        -7, -6, 2, 6, 9, 10, 12, 13, 14, 17
    };

    int sarray3[10]{
        -3, -2, -1, 0, 1, 2, 4, 8, 13, 16
    };

    std::cout << "These should all return 10:"
              << std::endl;

    std::cout << equal_to_index( sarray1, 10 )
              << std::endl;

    std::cout << equal_to_index( sarray2, 10 )
              << std::endl;

    std::cout << equal_to_index( sarray3, 10 )
              << std::endl;

    assert( equal_to_index( sarray1, 10 ) == 10 );
    assert( equal_to_index( sarray2, 10 ) == 10 );
    assert( equal_to_index( sarray3, 10 ) == 10 );

    sarray1[2] = 2;
    sarray2[3] = 3;
    sarray3[7] = 7;

    std::cout << std::endl
              << "These should return 2, 3, and 7:"
              << std::endl;

    std::cout << equal_to_index( sarray1, 10 )
              << std::endl;

    std::cout << equal_to_index( sarray2, 10 )
              << std::endl;

    std::cout << equal_to_index( sarray3, 10 )
              << std::endl;

    assert( equal_to_index( sarray1, 10 ) == 2 );
    assert( equal_to_index( sarray2, 10 ) == 3 );
    assert( equal_to_index( sarray3, 10 ) == 7 );

    // Unsorted arrays. Entries may be repeated.
    int uarray1[10]{
        16, 12, -1, 18, 9, -4, 4, 12, 13, 6
    };

    int uarray2[10]{
        19, 5, 7, 0, -2, -7, -5, -8, 2, -1
    };

    int uarray3[10]{
        8, 4, 16, 3, 12, -5, -5, -3, -1, 11
    };

    std::cout << std::endl
              << "These should all return false (0):"
              << std::endl;

    std::cout << equal_to_index( 2, uarray1, 10 )
              << std::endl;

    std::cout << equal_to_index( 8, uarray2, 10 )
              << std::endl;

    std::cout << equal_to_index( 6, uarray3, 10 )
              << std::endl;

    assert( !equal_to_index( 2, uarray1, 10 ) );
    assert( !equal_to_index( 8, uarray2, 10 ) );
    assert( !equal_to_index( 6, uarray3, 10 ) );

    uarray1[6] = 2;
    uarray2[2] = 8;
    uarray3[0] = 6;

    std::cout << std::endl
              << "These should all return true (1):"
              << std::endl;

    std::cout << equal_to_index( 2, uarray1, 10 )
              << std::endl;

    std::cout << equal_to_index( 8, uarray2, 10 )
              << std::endl;

    std::cout << equal_to_index( 6, uarray3, 10 )
              << std::endl;

    assert( equal_to_index( 2, uarray1, 10 ) );
    assert( equal_to_index( 8, uarray2, 10 ) );
    assert( equal_to_index( 6, uarray3, 10 ) );

    /*********************************************************
     * Additional tests for Problem 1
     *********************************************************/

    // More than one matching index:
    // return the smallest.
    int multiple[]{
        -3, 0, 2, 3, 5, 7
    };

    assert( equal_to_index( multiple, 6 ) == 2 );

    // Every entry matches its index.
    int all_match[]{
        0, 1, 2, 3, 4
    };

    assert( equal_to_index( all_match, 5 ) == 0 );

    // No entries match.
    int no_match[]{
        1, 2, 3, 4, 5
    };

    assert( equal_to_index( no_match, 5 ) == 5 );

    // Empty array.
    assert( equal_to_index( nullptr, 0 ) == 0 );

    /*********************************************************
     * Additional tests for Problem 2
     *********************************************************/

    // After sorting, the array is
    //
    //     {0, 1, 2, 2, 2, 5}.
    //
    // Both indices 2 and 3 contain 2, but only
    // the entry at index 2 equals its index.
    int duplicates[]{
        2, 5, 2, 0, 2, 1
    };

    assert( equal_to_index( 2, duplicates, 6 ) );
    assert( !equal_to_index( 3, duplicates, 6 ) );

    // Every entry is equal to 3.
    int identical[]{
        3, 3, 3, 3, 3
    };

    assert( equal_to_index( 3, identical, 5 ) );
    assert( !equal_to_index( 2, identical, 5 ) );

    return 0;
}

/*************************************************************
 * PROBLEM 1: MODIFIED BINARY SEARCH
 *
 * Precondition:
 *   The array contains distinct integers in
 *   strictly increasing order.
 *
 * Return:
 *   The smallest index k such that array[k] == k,
 *   or capacity if no such index exists.
 *
 * Define
 *
 *     f(k) = array[k] - k.
 *
 * Since the array contains distinct integers in
 * increasing order, f(k) is nondecreasing.
 *
 * Therefore:
 *
 *   If array[mid] < mid, no index at or before
 *   mid can satisfy array[k] == k.
 *
 *   Otherwise, a solution might exist at mid
 *   or to its left.
 *
 * We find the first index where array[k] >= k,
 * then check whether equality holds there.
 *
 * Time complexity:  O(log capacity)
 * Extra space:      O(1)
 *************************************************************/

std::size_t equal_to_index(
    int const         array[],
    std::size_t const capacity
) {
    assert( array != nullptr || capacity == 0 );

    std::size_t left{ 0 };
    std::size_t right{ capacity };

    // Maintain the half-open search interval
    // [left, right).
    while ( left < right ) {
        std::size_t const mid{
            left + (right - left)/2
        };

        // Compare using a signed representation of
        // the index. The exercise assumes the array
        // capacity does not exceed the largest int.
        assert(
            mid <= static_cast<std::size_t>(
                std::numeric_limits<int>::max()
            )
        );

        if ( array[mid] < static_cast<int>( mid ) ) {
            left = mid + 1;
        } else {
            right = mid;
        }
    }

    // left is the first index for which
    // array[left] >= left, if one exists.
    if (
        left < capacity
        && array[left] == static_cast<int>( left )
    ) {
        return left;
    }

    return capacity;
}

/*************************************************************
 * PROBLEM 2: COUNTING INSTEAD OF SORTING
 *
 * Precondition:
 *   idx < capacity.
 *
 * We want to know whether the entry at index idx
 * in the sorted array would equal idx.
 *
 * Let:
 *
 *   less  = number of entries strictly less than idx
 *
 *   equal = number of entries equal to idx
 *
 * After sorting, all entries equal to idx occupy
 * the half-open interval
 *
 *     [less, less + equal).
 *
 * Therefore, the answer is true exactly when
 *
 *     less <= idx < less + equal.
 *
 * This works even when entries are repeated.
 *
 * We never need to sort or modify the array.
 *
 * Time complexity:  O(capacity)
 * Extra space:      O(1)
 *************************************************************/

bool equal_to_index(
    std::size_t const idx,
    int const         array[],
    std::size_t const capacity
) {
    assert( array != nullptr );
    assert( idx < capacity );

    // As in Problem 1, indices must fit
    // in the integer type used by the array.
    assert(
        idx <= static_cast<std::size_t>(
            std::numeric_limits<int>::max()
        )
    );

    int const target{ static_cast<int>( idx ) };

    std::size_t less{ 0 };
    std::size_t equal{ 0 };

    for (
        std::size_t k{ 0 };
        k < capacity;
        ++k
    ) {
        if ( array[k] < target ) {
            ++less;
        } else if ( array[k] == target ) {
            ++equal;
        }
    }

    return less <= idx && idx - less < equal;
}