[leetcode] problem 863 – all nodes distance k in binary tree

public class  {
    int val;
    TreeNode left;
    TreeNode right;
    TreeNode(int x) { val = x; }
}

public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
    List<Integer> result = new ArrayList<>();
    helper(result, root, target, K);
    return result;
}

private int helper(List<Integer> result, TreeNode root, TreeNode target, int K) {
    if (root == null)
        return -1;

    if (root == target) {
        collect(result, root, K);
        return 0;
    }

    int left = helper(result, root.left, target, K);
    int right = helper(result, root.right, target, K);

    if (left >= 0) {
        if (left + 1 == K)
            result.add(root.val);

        collect(result, root.right, K - left - 2);
        return left + 1;
    }
    else if (right >= 0) {
        if (right + 1 == K)
            result.add(root.val);

        collect(result, root.left, K - right - 2);
        return right + 1;
    }
    else
        return -1;
}

private void collect(List<Integer> result, TreeNode root, int K) {
    if (root == null || K < 0)
        return;

    if (K == 0) {
        result.add(root.val);
        return;
    }

    collect(result, root.left, K - 1);
    collect(result, root.right, K - 1);
}