public class Kingdom {
    private final int n;              // number of villages
    private int numberOfProvinces;    // number of provinces
    private final int[] banners;      // banners[i] = j if village i flies banner j

    // creates a new kingdom with n villages, each initially in its own province
    public Kingdom(int n) {
        this.n = n;
        this.banners = new int[n];
        for (int i = 0; i < n; i++) {
            banners[i] = i;
        }
        numberOfProvinces = n;
    }

    // returns a string representation of this Kingdom
    public String toString() {
        String s = "[";
        for (int i = 0; i < banners.length; i++) {
            if (i > 0) {
                s += ", ";
            }
            s += banners[i];
        }
        return s + "]";
    }

    // returns the banner that village p flies
    public int bannerOf(int p) {
        checkIsValidVillage(p);
        return banners[p];
    }

    // returns the number of provinces
    public int provinceCount() {
        return numberOfProvinces;
    }

    // returns true if p and q are in the same province; false otherwise
    public boolean inSameProvince(int p, int q) {
        return bannerOf(p) == bannerOf(q);
    }

    // returns the number of villages in the largest province
    public int largestProvinceSize() {
        int[] sizes = new int[n];
        for (int i = 0; i < n; i++) {
            sizes[banners[i]]++;
        }
        int max = 0;
        for (int i = 0; i < n; i++) {
            if (sizes[i] > max) max = sizes[i];
        }
        return max;
    }

    // merges the province containing village p into the province containing village q
    public void signTreaty(int p, int q) {
        int bannerP = bannerOf(p);
        int bannerQ = bannerOf(q);
        if (bannerP == bannerQ) return;

        numberOfProvinces--;

        for (int i = 0; i < n; i++) {
            if (banners[i] == bannerP)
                banners[i] = bannerQ;
        }
    }

    // throw an exception if p is outside prescribed range (between 0 and n-1)
    private void checkIsValidVillage(int p) {
        if (p < 0 || p >= n) {
            throw new IllegalArgumentException("invalid village: " + p);
        }
    }


    /*************************************************************************
     *  unit tests the Kingdom data type
     *  the expected output is:
     *  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
     *  [6, 4, 7, 7, 4, 4, 6, 7, 7, 6]
     *  bannerOf(2)           = 7
     *  bannerOf(5)           = 4
     *  inSameProvince(2, 5)  = false
     *  provinceCount()       = 3
     *  largestProvinceSize() = 4
     *  [6, 4, 4, 4, 4, 4, 6, 4, 4, 6]
     *************************************************************************/
    public static void main(String[] args) {
        Kingdom kingdom = new Kingdom(10);
        StdOut.println(kingdom);
        kingdom.signTreaty(2, 3);
        kingdom.signTreaty(8, 7);
        kingdom.signTreaty(3, 7);
        kingdom.signTreaty(0, 9);
        kingdom.signTreaty(5, 4);
        kingdom.signTreaty(0, 6);
        kingdom.signTreaty(1, 4);
        StdOut.println(kingdom);
        StdOut.println("bannerOf(2)           = " + kingdom.bannerOf(2));
        StdOut.println("bannerOf(5)           = " + kingdom.bannerOf(5));
        StdOut.println("inSameProvince(2, 5)  = " + kingdom.inSameProvince(2, 5));
        StdOut.println("provinceCount()       = " + kingdom.provinceCount());
        StdOut.println("largestProvinceSize() = " + kingdom.largestProvinceSize());
        kingdom.signTreaty(3, 1);
        StdOut.println(kingdom);
    }
}
