package leetcode;

/**
 * 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;
        }
    }

    public int rob(TreeNode root) {
        // TODO
        int a = rob(root, false);
        int b = rob(root, true);
        System.out.println(a + " vs " + b);
        return Math.max(a, b);
    }

    private int rob(TreeNode node, boolean skip) {
        if (node == null) {
            return 0;
        }
        int val = 0;
        if (!skip) {
            val = node.val;
        }
        int a = rob(node.left, !skip);
        int b = rob(node.right, !skip);
        return a + b + val;
    }

    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
    }
}

@@ -12,16 +12,65 @@
 tcode;%0A%0A
+import java.util.HashMap;%0Aimport java.util.Map;%0A%0A
 /**%0A * h
@@ -358,15 +358,53 @@
     
-// TODO
+Map%3CInteger, Integer%3E memo = new HashMap%3C%3E();
 %0A   
@@ -430,13 +430,21 @@
 ot, 
-false
+0, 0, 0, memo
 );%0A 
@@ -451,24 +451,25 @@
        int b
+1
  = rob(root,
@@ -471,59 +471,103 @@
 root
-, true);%0A        System.ou
+.left, 0, 0, 0, memo);%0A        int b2 = rob(roo
 t.
-p
 ri
-ntln(a + %22 vs %22
+ght, 0, 0, 0, memo);%0A        int b = b1
  + b
-)
+2
 ;%0A  
@@ -641,20 +641,69 @@
 de, 
-boolean skip
+int idx, int step, int targetStep, Map%3CInteger, Integer%3E memo
 ) %7B%0A
@@ -796,157 +796,614 @@
    i
-f (!skip) %7B%0A            val = node.val;%0A        %7D%0A        int a = rob(node.left, !skip);%0A        int b = rob(node.right, !skip);%0A        return a + b
+nt newStep = step;%0A        if (step == targetStep) %7B%0A            val = node.val;%0A            newStep = 0;%0A        %7D%0A        int leftIdx = idx + 1;%0A        int left1 = rob(node.left, leftIdx, newStep + 1, 2, memo);%0A        int left2 = rob(node.left, leftIdx, newStep + 1, 3, memo);%0A        int maxLeft = Math.max(left1, left2);%0A%0A        int rightIdx = idx + 2;%0A        int right1 = rob(node.right, rightIdx, newStep + 1, 2, memo);%0A        int right2 = rob(node.right, rightIdx, newStep + 1, 3, memo);%0A        int maxRight = Math.max(right1, right2);%0A%0A        int max = maxLeft + maxRight;%0A        return max
  + v
@@ -2248,16 +2248,255 @@
 )); // 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    %7D%0A%7D

Update problem 337