findLake.java

public String findLake(int[] terrain, int w, int p) {
        int[] heights = flowWater(terrain, w, p);
        int maxHeight = heights[0];
        for (int i = 1; i < heights.length; i++) {
            maxHeight = Math.max(maxHeight, heights[i]);
        }
        StringBuilder sb = new StringBuilder();
        for (int i = maxHeight; i >= 0; i--) {
            for (int j = 0; j < heights.length; j++) {
                if (i > heights[j]) {
                    sb.append(" ");
                } else if (i <= heights[j] && i > terrain[j]) {
                    sb.append("W");
                } else  {
                    sb.append("+");
                }
            }
            sb.append("\n");
        }
        return sb.toString();
    }

    public int[] flowWater(int[] terrain, int w, int p) {
        int len = terrain.length;
        Deque<Integer> leftStack = new ArrayDeque<>();
        leftStack.addLast(p);
        int left = p;
        Deque<Integer> rightStack = new ArrayDeque<>();
        rightStack.addLast(p);
        int right = p;
        int[] heights = Arrays.copyOf(terrain, terrain.length);
        for (int j = 0; j < w; j++) {
            // update left if needed
            for (int i = left - 1; i >= 0 && heights[i] <= heights[i + 1]; i--) {
                if (heights[i] < heights[leftStack.peekLast()]) {
                    leftStack.addLast(i);
                }
                left = i;
            }
            // update right if needed
            for (int i = right + 1; i < len && heights[i] <= heights[i - 1]; i++) {
                if (heights[i] < heights[rightStack.peekLast()]) {
                    rightStack.addLast(i);
                }
                right = i;
            }

            if (leftStack.peekLast() != p) {
                int pos = leftStack.pollLast();
                heights[pos] += 1;
                if (heights[pos] < heights[leftStack.peekLast()]) {
                    leftStack.addLast(pos);
                }
                if (left < pos) {
                    leftStack.addLast(pos - 1);
                }

            } else if (rightStack.peekLast() != p) {
                int pos = rightStack.pollLast();
                heights[pos] += 1;
                if (heights[pos] < heights[rightStack.peekLast()]) {
                    rightStack.addLast(pos);
                }
                if (right > pos) {
                    rightStack.addLast(pos + 1);
                }
            } else {
                heights[p] += 1;
                if (left < p && p - 1 >= 0) {
                    leftStack.addLast(p - 1);
                }
                if (right > p && p + 1 < len) {
                    rightStack.addLast(p + 1);
                }
            }
        }
        return heights;
    }

FindLakes.java

import org.junit.Test;

import java.io.*;
import java.util.*;

public class FindLakes {
    class InputValues {
        int n;
        int[] terrain;
        int w;
        int p;

        public InputValues(int n, int[] terrain, int w, int p) {
            this.n = n;
            this.terrain = terrain;
            this.w = w;
            this.p = p;
        }
        @Override
        public String toString(){
            return n + "\n" + Arrays.toString(terrain) + "\n" + w + "\t" + p;
        }
    }

    @Test
    public void test() {
        File[] files = readFolder("src/input/FindLakes");
        FindLakes.InputValues test = readFile(files[1]);
        findLake(test.terrain, test.w, test.p);
    }
    
    public void generateOutput() throws IOException {
        File[] files = readFolder("src/input/FindLakes");
        for (int i = 0; i < files.length; i++) {
            File file = files[i];
            String outfileName = "src/output/FindLakes/" + file.getName().replace(".in", ".out");
            File outfile = new File(outfileName);
            outfile.getParentFile().mkdirs();
            outfile.createNewFile();
            PrintWriter printWriter = new PrintWriter(new FileWriter(outfileName));
            FindLakes.InputValues test = readFile(file);
            String lakeResult = findLake(test.terrain, test.w, test.p);
            printWriter.print(lakeResult);
            printWriter.close();
        }
    }
    
    public String findLake(int[] terrain, int w, int p) {
        int[] heights = flowWater(terrain, w, p);
        int maxHeight = heights[0];
        for (int i = 1; i < heights.length; i++) {
            maxHeight = Math.max(maxHeight, heights[i]);
        }
        StringBuilder sb = new StringBuilder();
        for (int i = maxHeight; i >= 0; i--) {
            for (int j = 0; j < heights.length; j++) {
                if (i > heights[j]) {
                    sb.append(" ");
                } else if (i <= heights[j] && i > terrain[j]) {
                    sb.append("W");
                } else  {
                    sb.append("+");
                }
            }
            sb.append("\n");
        }
        return sb.toString();
    }

    public int[] flowWater(int[] terrain, int w, int p) {
        int len = terrain.length;
        Deque<Integer> leftStack = new ArrayDeque<>();
        leftStack.addLast(p);
        int left = p;
        Deque<Integer> rightStack = new ArrayDeque<>();
        rightStack.addLast(p);
        int right = p;
        int[] heights = Arrays.copyOf(terrain, terrain.length);
        for (int j = 0; j < w; j++) {
            // update left if needed
            for (int i = left - 1; i >= 0 && heights[i] <= heights[i + 1]; i--) {
                if (heights[i] < heights[leftStack.peekLast()]) {
                    leftStack.addLast(i);
                }
                left = i;
            }
            // update right if needed
            for (int i = right + 1; i < len && heights[i] <= heights[i - 1]; i++) {
                if (heights[i] < heights[rightStack.peekLast()]) {
                    rightStack.addLast(i);
                }
                right = i;
            }

            if (leftStack.peekLast() != p) {
                int pos = leftStack.pollLast();
                heights[pos] += 1;
                if (heights[pos] < heights[leftStack.peekLast()]) {
                    leftStack.addLast(pos);
                }
                if (left < pos) {
                    leftStack.addLast(pos - 1);
                }

            } else if (rightStack.peekLast() != p) {
                int pos = rightStack.pollLast();
                heights[pos] += 1;
                if (heights[pos] < heights[rightStack.peekLast()]) {
                    rightStack.addLast(pos);
                }
                if (right > pos) {
                    rightStack.addLast(pos + 1);
                }
            } else {
                heights[p] += 1;
                if (left < p && p - 1 >= 0) {
                    leftStack.addLast(p - 1);
                }
                if (right > p && p + 1 < len) {
                    rightStack.addLast(p + 1);
                }
            }
        }
        return heights;
    }

    private int getInt(String name) {
        return Integer.parseInt(name.replaceAll("\\D", ""));
    }

    private File[] readFolder(String path) {
        File folder = new File(path);
        File[] files = folder.listFiles();
        Arrays.sort(files, (a, b) -> getInt(a.getName()) - getInt(b.getName()));
        return files;
    }

    private FindLakes.InputValues readFile(File file) {
        try {
            // assume the input file is always valid
            Scanner sc = new Scanner(file);
            int n = Integer.parseInt(sc.nextLine().trim());
            String[] altitudesStr = sc.nextLine().split(" ");
            int[] altitude = new int[altitudesStr.length];
            for (int i = 0; i < altitude.length; i++) {
                altitude[i] = Integer.parseInt(altitudesStr[i]);
            }
            String[] pw = sc.nextLine().split(" ");
            int p = Integer.parseInt(pw[0]);
            int w = Integer.parseInt(pw[1]);
            sc.close();
            return new FindLakes.InputValues(n, altitude, p, w);
        } catch (FileNotFoundException e) {
            System.out.println(file.getName() + " not exist");
        }
        return null;
    }
  
}