Google FooBar challenge 5-1 Reverse Game of Life one step

Solution.java

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.StringTokenizer;
import java.util.stream.Collectors;
import java.util.stream.Stream;

public clas Solution {
      public static class Pair {
        private final int x;
        private final int y;

        public Pair(int x, int y) {
            this.x = x;
            this.y = y;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            Pair pair = (Pair) o;
            return x == pair.x && y == pair.y;
        }

        @Override
        public int hashCode() {
            return Objects.hash(x, y);
        }

        @Override
        public String toString() {
            return "Pair{" + "x=" + x + ", y=" + y + '}';
        }
    }

    private static boolean[][] toGrid(int state, int h, int w) {
        boolean[][] grid = new boolean[h][w];
        for (int i = 0; i < h; i++) {
            for (int j = 0; j < w; j++) {
                grid[i][j] = (state >> (i * w + j)) % 2 == 1;
            }
        }
        return grid;
    }

    private static int colState(boolean[][] g, int col) {
        if (g[0].length <= col) return 0;
        int s = 0;
        for (int row = 0; row < g.length; row++) {
            if (g[row][col]) {
                s += (1 << row);
            }
        }
        return s;
    }

    // list of all possible states (list of 2x2), if the state is 1
    private static final boolean[][][] POSITIVE = Stream
            .of(1, 2, 4, 8)
            .map(s -> toGrid(s, 2, 2))
            .toArray(boolean[][][]::new);

    // list of all possible states (list of 2x2), if the state is 0
    private static final boolean[][][] NEGATIVE = Stream
            .of(0, 3, 5, 6, 7, 9, 10, 11, 12, 13, 14, 15)
            .map(s -> toGrid(s, 2, 2))
            .toArray(boolean[][][]::new);

    // current number state (hx1, 512)
    //  -> list of all possible previous states (list of (h+1)x2)
    private static Map<Integer, List<Pair>> PREV_STATES;

    // current number state (hx1, 512)
    //  -> left side prev number state ((h+1)x1, 1024)
    //  -> list of all possible previous states (right side) (list of (h+1)x1)
    private static Map<Integer, Map<Integer, List<Integer>>> PREV_FITTING_STATES;

    private static void preCalcStates(Set<Integer> states, int h) {
        PREV_STATES = states.stream()
                .collect(Collectors.toMap(
                        s -> s,
                        s -> allPreviousStates(toGrid(s, h, 1))
                                .map(g -> new Pair(colState(g, 0), colState(g, 1)))
                                .collect(Collectors.toList())
                ));
        PREV_FITTING_STATES = PREV_STATES.entrySet().stream()
                .collect(Collectors.toMap(
                        Map.Entry::getKey,
                        e -> {
                            List<Pair> possibleStates = e.getValue();
                            Map<Integer, List<Integer>> fittingStates = new HashMap<>();
                            for (Pair s : possibleStates) {
                                if (!fittingStates.containsKey(s.x)) {
                                    fittingStates.put(s.x, new ArrayList<>());
                                }
                                fittingStates.get(s.x).add(s.y);
                            }
                            return fittingStates;
                        }
                ));
    }

    private static int ans = 0;

    // g[h][w]
    // 3 <= h <= 9
    // 3 <= w <= 50
    public static int solution(boolean[][] g) {
        ans = 0;

        int h = g.length;
        int w = g[0].length;
        int[] states = new int[w];
        Set<Integer> statesSet = new HashSet<>();
        for (int col = 0; col < g[0].length; col++) {
            states[col] = colState(g, col);
            statesSet.add(states[col]);
        }
        preCalcStates(statesSet, h);

        int[] cur = new int[w+1];
        checkRecBulk(states, 0, cur);
        return ans;
    }

    private static boolean[][] copy(boolean[][] src) {
        boolean[][] dst = new boolean[src.length][src[0].length];
        for (int i = 0; i < src.length; i++) {
            dst[i] = Arrays.copyOf(src[i], src[i].length);
        }
        return dst;
    }

    private static Stream<boolean[][]> allPreviousStates(boolean[][] g) {
        return allPreviousStatesI(g, 0, 0, new boolean[g.length+1][g[0].length+1]);
    }
    
    // naive solution
    private static Stream<boolean[][]> allPreviousStatesI(boolean[][] g, int i, int j, boolean[][] cur) {
        if (g.length <= i) {
            // because of 'of' method overloading we have to make one elt boolean[][][] array
            return Stream.of(new boolean[][][]{copy(cur)});
        }

        boolean isRowsLast = j == g[0].length - 1;
        boolean[][][] possibles = g[i][j] ? POSITIVE : NEGATIVE;
        return Arrays.stream(possibles).filter(possible -> {
            // check validity
            if ((i > 0 || j > 0) && possible[0][0] != cur[i][j]) return false;
            if (i > 0 && possible[0][1] != cur[i][j+1]) return false;
            if (j > 0 && possible[1][0] != cur[i+1][j]) return false;
            return true;
        }).flatMap(possible -> {
            // fill in and check next
            cur[i][j] = possible[0][0];
            cur[i][j+1] = possible[0][1];
            cur[i+1][j] = possible[1][0];
            cur[i+1][j+1] = possible[1][1];
            if (isRowsLast) {
                return allPreviousStatesI(g, i+1, 0, cur);
            } else {
                return allPreviousStatesI(g, i, j+1, cur);
            }
        });
    }

    // check in hx1 bulk
    private static void checkRecBulk(int[] g, int col, int[] cur) {
        final boolean isFirst = col == 0;
        final boolean isLast = col == g.length - 1;
        if (isFirst && isLast) {
            ans += PREV_STATES.get(g[col]).size();
            return;
        }
        if (isFirst) {
            List<Pair> states = PREV_STATES.get(g[col]);
            for (Pair state : states) {
                // fill in and check next
                cur[col] = state.x;
                cur[col+1] = state.y;
                checkRecBulk(g, col+1, cur);
            }
            return;
        }

        List<Integer> states = PREV_FITTING_STATES.get(g[col]).getOrDefault(cur[col], Collections.emptyList());
        if (isLast) {
            ans += states.size();
            return;
        }
        for (int state : states) {
            // fill in and check next
            cur[col+1] = state;
            checkRecBulk(g, col+1, cur);
        }
    }
}