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

/*************************************************************
 * DISJOINT SETS USING PARENT INDICES
 *
 * In the previous implementation of disjoint sets, two
 * entries m and n were in the same set exactly when
 *
 *     array[m] == array[n].
 *
 * Unfortunately, taking the union of two sets could require
 * searching the entire array and changing many entries.
 * Consider an image with a billion pixels. If two groups of
 * pixels are both identified as asphalt, joining those
 * groups should not require examining every pixel.
 *
 * Here is a different way to represent disjoint sets.
 *
 * Each entry stores the index of its parent. A set is
 * represented by a rooted tree, and its root identifies
 * the set. If array[m] == m, then m is a root. Otherwise,
 * follow parent indices until reaching a root:
 *
 *     unsigned int pm{ m };
 *
 *     while ( pm != array[pm] ) {
 *         pm = array[pm];
 *     }
 *
 * Now pm is the representative of the set containing m.
 *
 * To determine whether m and n are in the same set, find
 * their respective roots pm and pn and compare them.
 *
 * To join two different sets, find their roots and make
 * one root a child of the other:
 *
 *     array[pm] = pn;
 *
 * or
 *
 *     array[pn] = pm;
 *
 * Either assignment is correct if pm != pn. Which might
 * lead to shorter paths when finding roots later?
 *
 * EXERCISE
 *
 * 1. Initialize each entry as its own parent.
 * 2. Implement in_same_set() by following parent indices.
 * 3. Implement take_union() by joining two roots.
 * 4. Modify print_sets() to group entries by their roots,
 *    not by their immediate parents.
 *
 * Complete the functions below before reading the hints
 * and the commented solutions at the end of this file.
 *************************************************************/

// Function declarations
int main();

void init_sets( unsigned int       array[],
                std::size_t  const capacity );

bool in_same_set( unsigned int const m,
                  unsigned int const n,
                  unsigned int       array[],
                  std::size_t  const capacity );

void take_union( unsigned int const m,
                 unsigned int const n,
                 unsigned int       array[],
                 std::size_t  const capacity );

void print_sets( unsigned int const array[],
                 std::size_t  const capacity );

// Function definitions
int main() {
    unsigned int const N{ 10 };
    unsigned int array[N]{};

    init_sets( array, N );

    // {0}, {1}, {2}, {3}, {4}, {5}, {6}, {7}, {8}, {9}
    print_sets( array, N );

    take_union( 3, 5, array, N );
    take_union( 3, 6, array, N );
    take_union( 5, 7, array, N );

    // {0}, {1}, {2}, {3, 5, 6, 7}, {4}, {8}, {9}
    print_sets( array, N );

    assert( in_same_set( 3, 5, array, N ) );
    assert( in_same_set( 3, 6, array, N ) );
    assert( in_same_set( 5, 7, array, N ) );
    assert( !in_same_set( 3, 4, array, N ) );

    take_union( 0, 8, array, N );
    take_union( 1, 9, array, N );

    // {0, 8}, {1, 9}, {2}, {3, 5, 6, 7}, {4}
    print_sets( array, N );

    take_union( 9, 8, array, N );

    // {0, 1, 8, 9}, {2}, {3, 5, 6, 7}, {4}
    print_sets( array, N );

    assert( in_same_set( 0, 1, array, N ) );
    assert( in_same_set( 0, 8, array, N ) );
    assert( in_same_set( 0, 9, array, N ) );
    assert( !in_same_set( 0, 3, array, N ) );

    // Joining an already joined pair should change nothing.
    take_union( 0, 9, array, N );
    assert( in_same_set( 0, 9, array, N ) );

    return 0;
}

// Initialize the array so that every entry is its own root.
void init_sets( unsigned int       array[],
                std::size_t  const capacity ) {
    assert( array != nullptr || capacity == 0 );

    for ( std::size_t k{ 0 }; k < capacity; ++k ) {
        array[k] = static_cast<unsigned int>( k );
    }
}

// Return true if m and n have the same root.
bool in_same_set( unsigned int const m,
                  unsigned int const n,
                  unsigned int       array[],
                  std::size_t  const capacity ) {
    assert( array != nullptr );
    assert( m < capacity );
    assert( n < capacity );

    // Find the root of m and the root of n.
    // Compare those roots, not array[m] and array[n].
    // Enter your implementation here.
    return false;
}

// Join the sets containing m and n by linking their roots.
void take_union( unsigned int const m,
                 unsigned int const n,
                 unsigned int       array[],
                 std::size_t  const capacity ) {
    assert( array != nullptr );
    assert( m < capacity );
    assert( n < capacity );

    // Find both roots. If they are different, make one
    // root point to the other. Never link arbitrary nodes.
    // Enter your implementation here.
}

// Print the sets, for example:
// {0, 1}, {2, 3}, {4, 5}
// The order of the sets does not matter.
void print_sets( unsigned int const array[],
                 std::size_t  const capacity ) {
    assert( array != nullptr || capacity == 0 );

    bool first_set_printed{ false };
    std::vector<bool> entry_printed( capacity, false );

    for ( std::size_t k{ 0 }; k < capacity; ++k ) {
        if ( !entry_printed[k] ) {
            if ( first_set_printed ) {
                std::cout << ", ";
            } else {
                first_set_printed = true;
            }

            entry_printed[k] = true;
            std::cout << "{" << k;

            // Find the root of k. For each subsequent entry
            // i, find its root. Print i if the roots agree,
            // and mark i as printed.
            // Enter your implementation here.

            std::cout << "}";
        }
    }

    std::cout << std::endl;
}

/*************************************************************
 * HINTS
 *
 * Hint 1: Finding a root
 *
 * Start at m. While array[m] != m, replace m with array[m].
 * Be careful to update the variable you are following;
 * otherwise, the loop will never advance.
 *
 * Hint 2: Membership
 *
 * Find the roots pm and pn of m and n, respectively.
 * Return whether pm == pn. Two entries may have different
 * immediate parents yet still belong to the same set.
 *
 * Hint 3: Union
 *
 * Find pm and pn. If they differ, assign array[pm] = pn
 * or array[pn] = pm. Do not change array[m] or array[n]
 * unless those entries happen to be roots.
 *
 * Hint 4: Printing
 *
 * Retain the entry_printed array from the previous
 * exercise. Before comparing k and i, follow their
 * parent indices to their roots. Compare those roots.
 *
 * Hint 5: Choosing a parent
 *
 * Either root can become the parent of the other.
 * However, attaching a shorter tree beneath a taller
 * tree avoids unnecessarily increasing the height.
 * How could you keep track of tree heights? This is
 * an optional extension, not required for this exercise.
 *
 * If your code is not working, draw the parent pointers
 * after each union and follow them with a pencil.
 *************************************************************/

/*************************************************************
 * INSTRUCTOR'S SOLUTIONS
 *
 * These are replacement definitions for the three
 * incomplete functions above. They are deliberately
 * enclosed in a block comment so that the student
 * exercise remains incomplete when compiled.
 *************************************************************/

/*
bool in_same_set( unsigned int const m,
                  unsigned int const n,
                  unsigned int       array[],
                  std::size_t  const capacity ) {
    assert( array != nullptr );
    assert( m < capacity );
    assert( n < capacity );

    unsigned int pm{ m };
    unsigned int pn{ n };

    while ( pm != array[pm] ) {
        assert( array[pm] < capacity );
        pm = array[pm];
    }

    while ( pn != array[pn] ) {
        assert( array[pn] < capacity );
        pn = array[pn];
    }

    return pm == pn;
}

void take_union( unsigned int const m,
                 unsigned int const n,
                 unsigned int       array[],
                 std::size_t  const capacity ) {
    assert( array != nullptr );
    assert( m < capacity );
    assert( n < capacity );

    unsigned int pm{ m };
    unsigned int pn{ n };

    while ( pm != array[pm] ) {
        assert( array[pm] < capacity );
        pm = array[pm];
    }

    while ( pn != array[pn] ) {
        assert( array[pn] < capacity );
        pn = array[pn];
    }

    if ( pm != pn ) {
        array[pn] = pm;
    }
}

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

    bool first_set_printed{ false };
    std::vector<bool> entry_printed( capacity, false );

    for ( std::size_t k{ 0 }; k < capacity; ++k ) {
        if ( !entry_printed[k] ) {
            if ( first_set_printed ) {
                std::cout << ", ";
            } else {
                first_set_printed = true;
            }

            entry_printed[k] = true;
            std::cout << "{" << k;

            unsigned int pk{ static_cast<unsigned int>( k ) };

            while ( pk != array[pk] ) {
                assert( array[pk] < capacity );
                pk = array[pk];
            }

            for ( std::size_t i{ k + 1 }; i < capacity; ++i ) {
                unsigned int pi{ static_cast<unsigned int>( i ) };

                while ( pi != array[pi] ) {
                    assert( array[pi] < capacity );
                    pi = array[pi];
                }

                if ( pi == pk ) {
                    std::cout << ", " << i;
                    entry_printed[i] = true;
                }
            }

            std::cout << "}";
        }
    }

    std::cout << std::endl;
}
*/