package leetcode;

import java.util.HashMap;
import java.util.Map;

/**
 * https://leetcode.com/problems/house-robber-iii/
 */
public class Problem337 {
    public static class TreeNode {
        int val;
        TreeNode left;
        TreeNode right;

        TreeNode(int x) {
            val = x;
        }

        @Override
        public String toString() {
            final StringBuilder sb = new StringBuilder("TreeNode{");
            sb.append("val=").append(val);
            sb.append('}');
            return sb.toString();
        }
    }

    public int rob(TreeNode root) {
        if (root == null) {
            return 0;
        }
        Map<TreeNodeWrapper, Integer> memo = new HashMap<>();
        int a = rob(root, false, 0, memo);
        int b = rob(root, true, 0, memo);
//        System.out.println("a: " + a + " vs b: " + b);
        return Math.max(a, b);
    }

    private static class TreeNodeWrapper {
        private TreeNode node;
        private boolean skipped;

        public TreeNodeWrapper(TreeNode node, boolean skipped) {
            this.node = node;
            this.skipped = skipped;
        }
    }

    private int rob(TreeNode node, boolean skipped, int idx, Map<TreeNodeWrapper, Integer> memo) {
        if (node == null) {
            return 0;
        }
        TreeNodeWrapper tnw = new TreeNodeWrapper(node, skipped);
        if (memo.containsKey(tnw)) {
            return memo.get(tnw);
        }
        int max = 0;
        int leftIdx = idx + 1;
        int rightIdx = idx + 2;
        if (skipped) {
            int left1 = rob(node.left, skipped, leftIdx, memo); // continue skipping
            int left2 = rob(node.left, !skipped, leftIdx, memo); // start processing
            int maxLeft = Math.max(left1, left2);

            int right1 = rob(node.right, skipped, rightIdx, memo); // continue skipping
            int right2 = rob(node.right, !skipped, rightIdx, memo); // start processing
            int maxRight = Math.max(right1, right2);

            max = maxLeft + maxRight;
        } else { // not skipped
            int maxLeft = rob(node.left, !skipped, leftIdx, memo); // start skipping
            int maxRight = rob(node.right, !skipped, rightIdx, memo); // start skipping
            max = maxLeft + maxRight;
        }
        if (!skipped) {
            max += node.val;
        }
        memo.put(tnw, max);
        return max;
    }

    public static void main(String[] args) {
        Problem337 prob = new Problem337();

        TreeNode root = new TreeNode(3);
        root.left = new TreeNode(2);
        root.left.right = new TreeNode(3);
        root.right = new TreeNode(3);
        root.right.right = new TreeNode(1);
        System.out.println(prob.rob(root)); // 7

        root = new TreeNode(3);
        root.left = new TreeNode(4);
        root.left.left = new TreeNode(1);
        root.left.right = new TreeNode(3);
        root.right = new TreeNode(5);
        root.right.right = new TreeNode(1);
        System.out.println(prob.rob(root)); // 9

        root = new TreeNode(4);
        root.left = new TreeNode(1);
        root.left.left = new TreeNode(2);
        root.left.left.left = new TreeNode(3);
        System.out.println(prob.rob(root)); // 7

        root = new TreeNode(3);
        System.out.println(prob.rob(root)); // 3

        root = new TreeNode(3);
        root.left = new TreeNode(1);
        root.right = new TreeNode(3);
        System.out.println(prob.rob(root)); // 4

        root = new TreeNode(2);
        root.left = new TreeNode(1);
        root.left.right = new TreeNode(4);
        root.right = new TreeNode(3);
        System.out.println(prob.rob(root)); // 7
    }
}

@@ -59,16 +59,42 @@
 til.Map;
+%0Aimport java.util.Objects;
 %0A%0A/**%0A *
@@ -332,246 +332,8 @@
    %7D
-%0A%0A        @Override%0A        public String toString() %7B%0A            final StringBuilder sb = new StringBuilder(%22TreeNode%7B%22);%0A            sb.append(%22val=%22).append(val);%0A            sb.append('%7D');%0A            return sb.toString();%0A        %7D
 %0A   
@@ -526,19 +526,16 @@
 , false,
- 0,
  memo);%0A
@@ -569,19 +569,16 @@
 rue,
- 0,
  memo);%0A
 //  
@@ -577,65 +577,8 @@
 o);%0A
-//        System.out.println(%22a: %22 + a + %22 vs b: %22 + b);%0A
     
@@ -852,36 +852,478 @@
 pped;%0A        %7D%0A
+%0A
     
+    @Override%0A        public boolean equals(Object o) %7B%0A            if (this == o) return true;%0A            if (o == null %7C%7C getClass() != o.getClass()) return false;%0A            TreeNodeWrapper that = (TreeNodeWrapper) o;%0A            return skipped == that.skipped &&%0A                Objects.equals(node, that.node);%0A        %7D%0A%0A        @Override%0A        public int hashCode() %7B%0A            return Objects.hash(node, skipped);%0A        %7D%0A    
 %7D%0A%0A    private i
@@ -1363,17 +1363,8 @@
 pped
-, int idx
 , Ma
@@ -1631,71 +1631,8 @@
  0;%0A
-        int leftIdx = idx + 1;%0A        int rightIdx = idx + 2;%0A
     
@@ -1693,33 +1693,24 @@
 ft, skipped,
- leftIdx,
  memo); // c
@@ -1770,33 +1770,24 @@
 t, !skipped,
- leftIdx,
  memo); // s
@@ -1898,34 +1898,24 @@
 ht, skipped,
- rightIdx,
  memo); // c
@@ -1977,34 +1977,24 @@
 t, !skipped,
- rightIdx,
  memo); // s
@@ -2187,17 +2187,8 @@
 ped,
- leftIdx,
  mem
@@ -2265,18 +2265,8 @@
 ped,
- rightIdx,
  mem
@@ -2455,1289 +2455,7 @@
    %7D
-%0A%0A    public static void main(String%5B%5D args) %7B%0A        Problem337 prob = new Problem337();%0A%0A        TreeNode root = new TreeNode(3);%0A        root.left = new TreeNode(2);%0A        root.left.right = new TreeNode(3);%0A        root.right = new TreeNode(3);%0A        root.right.right = new TreeNode(1);%0A        System.out.println(prob.rob(root)); // 7%0A%0A        root = new TreeNode(3);%0A        root.left = new TreeNode(4);%0A        root.left.left = new TreeNode(1);%0A        root.left.right = new TreeNode(3);%0A        root.right = new TreeNode(5);%0A        root.right.right = new TreeNode(1);%0A        System.out.println(prob.rob(root)); // 9%0A%0A        root = new TreeNode(4);%0A        root.left = new TreeNode(1);%0A        root.left.left = new TreeNode(2);%0A        root.left.left.left = new TreeNode(3);%0A        System.out.println(prob.rob(root)); // 7%0A%0A        root = new TreeNode(3);%0A        System.out.println(prob.rob(root)); // 3%0A%0A        root = new TreeNode(3);%0A        root.left = new TreeNode(1);%0A        root.right = new TreeNode(3);%0A        System.out.println(prob.rob(root)); // 4%0A%0A        root = new TreeNode(2);%0A        root.left = new TreeNode(1);%0A        root.left.right = new TreeNode(4);%0A        root.right = new TreeNode(3);%0A        System.out.println(prob.rob(root)); // 7%0A    %7D
 %0A%7D%0A

Solve problem 337