gistfile1.java

import com.google.common.base.Joiner;

import java.io.*;
import java.util.*;
import static java.lang.System.out;
import static java.lang.String.format;

public class DraainageDraaaaainage {
    private final int[][] elevations;
    private final Node[][] nodes;

    private final int rows;
    private final int cols;

    private char basinLabel = 'a';

    public DraainageDraaaaainage(int[][] elevations) {
        this.elevations = elevations;

        this.rows = elevations.length;
        this.cols = elevations[0].length;

        this.nodes = new Node[rows][cols];
    }

    private static class Node {
        int row;
        int col;
        int elevation;

        Character basinLabel;

        Node outgoing;
        List<Node> incoming = new LinkedList<Node>();

        Node(int row, int col, int elevation) {
            this.row = row;
            this.col = col;
            this.elevation = elevation;
        }
    }

    private static class NodeComparator implements Comparator<Node> {
        private int pivotRow;
        private int pivotCol;

        NodeComparator(int pivotRow, int pivotCol) {
            this.pivotRow = pivotRow;
            this.pivotCol = pivotCol;
        }

        private int directionWeight(Node node) {
            int rowDiff = node.row - pivotRow;
            int colDiff = node.col - pivotCol;

            if (rowDiff == -1)
                return 1; // North
            else if (colDiff == -1)
                return 2; // West
            else if (colDiff == 1)
                return 3; // East
            else if (rowDiff == 1)
                return 4; // South
            else
                throw new IllegalArgumentException("Either row or col diff must be 1 or -1, with the other value being zero.");
        }

        public int compare(Node o1, Node o2) {
            int elevationDiff = o1.elevation - o2.elevation;

            if (elevationDiff != 0)
                return elevationDiff;
            else
                return directionWeight(o1) - directionWeight(o2);
        }
    }

    Character[][] getBasinLabels() {
        int rows = elevations.length;
        int cols = elevations[0].length;

        if (rows == 1 && cols == 1)
            return new Character[][] {new Character[] {'a'}};

        Node[][] nodes = new Node[rows][cols];

        // construct the matrix of nodes
        for (int i = 0; i < rows; i++)
            for (int j = 0; j < cols; j++)
                populateNode(i, j);

        // find the basins and their individual cell labels
        Character[][] basinLabels = new Character[rows][cols];
        for (int i = 0; i < rows; i++)
            for (int j = 0; j < cols; j++)
                basinLabels[i][j] = getBasinLabel(i, j);

        return basinLabels;
    }

    void populateNode(int i, int j) {
        int[][] adjacentCoords = new int[][]{
            new int[]{i, j - 1},
            new int[]{i, j + 1},
            new int[]{i - 1, j},
            new int[]{i + 1, j}
        };

        List<Node> adjacentNodes = new ArrayList<Node>(4);

        for (int[] coord : adjacentCoords) {
            if (coord[0] < 0 || coord[0] >= rows || coord[1] < 0 || coord[1] >= cols)
                continue;

            adjacentNodes.add(getOrCreateNode(coord[0], coord[1]));
        }

        Collections.sort(adjacentNodes, new NodeComparator(i, j));
        Node outgoing = adjacentNodes.get(0);

        Node node = getOrCreateNode(i, j);

        // otherwise this is a sink, and outgoing remains null
        if (outgoing.elevation < node.elevation) {
            node.outgoing = outgoing;
            node.outgoing.incoming.add(node);
        }

        nodes[i][j] = node;
    }

    Node getOrCreateNode(int i, int j) {
        if (nodes[i][j] == null)
            nodes[i][j] = new Node(i, j, elevations[i][j]);
        return nodes[i][j];
    }

    char getBasinLabel(int i, int j) {
        Node node = nodes[i][j];

        // a previous invocation already labeled this cell
        if (node.basinLabel != null)
            return node.basinLabel;

        Node cursor = node;
        while (cursor.outgoing != null)
            cursor = cursor.outgoing;

        Node sink = cursor;
        char label = basinLabel++;

        Stack<Node> toVisit = new Stack<Node>();
        toVisit.add(sink);

        // DFS to back-track through all cells in the basin, labeling them as we go
        while (toVisit.size() > 0) {
            cursor = toVisit.pop();
            cursor.basinLabel = label;

            for (Node in : cursor.incoming)
                toVisit.push(in);
        }

        return label;
    }

    public static void main(String[] args) throws IOException {
        BufferedReader reader = new BufferedReader(new FileReader(new File("input.txt")));

        int maps = Integer.parseInt(reader.readLine());

        for (int i = 0; i < maps; i++) {
            String[] header = reader.readLine().split(" ");
            int rows = Integer.parseInt(header[0]);
            int cols = Integer.parseInt(header[1]);

            int[][] elevations = new int[rows][cols];

            for (int j = 0; j < rows; j++) {
                String[] rowCells = reader.readLine().split(" ");

                for (int k = 0; k < cols; k++)
                    elevations[j][k] = Integer.parseInt(rowCells[k]);
            }

            DraainageDraaaaainage drainage = new DraainageDraaaaainage(elevations);
            Character[][] basinMap = drainage.getBasinLabels();

            out.println(format("Case #%d:", i + 1));

            Joiner joiner = Joiner.on(" ");
            for (Character[] row : basinMap)
                out.println(joiner.join(row));
        }
    }

    /*
    public static void main(String[] args) {
        int[][] input = new int[][]{
            new int[]{9, 6, 3},
            new int[]{5, 9, 6},
            new int[]{3, 5, 9}
        };

        char[][] expected = new char[][]{
            new char[]{'a', 'b', 'b'},
            new char[]{'a', 'a', 'b'},
            new char[]{'a', 'a', 'a'}
        };

        DraainageDraaaaainage drainage = new DraainageDraaaaainage(input);
        ReflectionAssert.assertReflectionEquals(expected, drainage.getBasinLabels());
    }
    */
}