863. All Nodes Distance K in Binary Tree

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    List<Integer> list = new ArrayList<>();
    int val;
    private void searchSubTree(TreeNode root, int count, int k) {
        if (root == null) return;
        if (count == k) {
            list.add(root.val);
            return;
        }
        searchSubTree(root.left, count + 1, k);
        searchSubTree(root.right, count + 1, k);
    }

    private int trav(TreeNode root, TreeNode target, int k, int count) {
        if (root == null) return 0;
        if (root == target) {
            searchSubTree(target, 0, k);
            return 1;
        }

        int newCount = count + 1;

        int left = trav(root.left, target, k, count > 0 ? newCount : count);
        int right = trav(root.right, target, k, left > 0 ? newCount : count);
        if (left > 0) searchSubTree(root.right, left + 1, k);
        else if (right > 0) searchSubTree(root.left, right + 1, k);

        if ((left == k || right == k) && k > 0) list.add(root.val);
        return left > 0 ? left + 1 : right > 0 ? right + 1 : 0;
    }
    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        trav(root, target, k, 0);
        return list;
    }
}