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<String, Integer> memo = new HashMap<>();
        int a = rob(root, false, memo);
        int b = rob(root, true, memo);
//        System.out.println("a: " + a + " vs b: " + b);
        return Math.max(a, b);
    }

    private int rob(TreeNode node, boolean skipped, Map<String, Integer> memo) {
        if (node == null) {
            return 0;
        }
        int max;
        if (skipped) {
            int left1 = rob(node.left, skipped, memo); // continue skipping
            int left2 = rob(node.left, !skipped, memo); // start processing
            int maxLeft = Math.max(left1, left2);

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

            max = maxLeft + maxRight;
        } else { // not skipped
            int maxLeft = rob(node.left, !skipped, memo); // start skipping
            int maxRight = rob(node.right, !skipped, memo); // start skipping
            max = maxLeft + maxRight;
        }
        if (!skipped) {
            max += node.val;
        }
        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
    }
}

@@ -652,30 +652,39 @@
         Map%3C
-String
+TreeNodeWrapper
 , Integer%3E m
@@ -738,16 +738,19 @@
 , false,
+ 0,
  memo);%0A
@@ -780,16 +780,19 @@
 t, true,
+ 0,
  memo);%0A
@@ -882,24 +882,280 @@
  b);%0A    %7D%0A%0A
+    private static class TreeNodeWrapper %7B%0A        private TreeNode node;%0A        private boolean skipped;%0A%0A        public TreeNodeWrapper(TreeNode node, boolean skipped) %7B%0A            this.node = node;%0A            this.skipped = skipped;%0A        %7D%0A    %7D%0A%0A
     private 
@@ -1198,18 +1198,36 @@
 ed, 
-Map%3CString
+int idx, Map%3CTreeNodeWrapper
 , In
@@ -1309,23 +1309,237 @@
         
-int max
+TreeNodeWrapper tnw = new TreeNodeWrapper(node, skipped);%0A        if (memo.containsKey(tnw)) %7B%0A            return memo.get(tnw);%0A        %7D%0A        int max = 0;%0A        int leftIdx = idx + 1;%0A        int rightIdx = idx + 2
 ;%0A      
@@ -1594,32 +1594,41 @@
 e.left, skipped,
+ leftIdx,
  memo); // conti
@@ -1680,32 +1680,41 @@
 .left, !skipped,
+ leftIdx,
  memo); // start
@@ -1817,32 +1817,42 @@
 .right, skipped,
+ rightIdx,
  memo); // conti
@@ -1906,32 +1906,42 @@
 right, !skipped,
+ rightIdx,
  memo); // start
@@ -2118,32 +2118,41 @@
 .left, !skipped,
+ leftIdx,
  memo); // start
@@ -2213,16 +2213,26 @@
 skipped,
+ rightIdx,
  memo); 
@@ -2352,32 +2352,60 @@
 .val;%0A        %7D%0A
+        memo.put(tnw, max);%0A
         return m

Update problem 337