java O(logn + k) use iterator, how can I simplify my code?


  • 0
    M

    here I create two iterators, first find the closest node in tree, then run k times next().....

    public List<Integer> closestKValues(TreeNode root, double target, int k) {

    	List<Integer> res = new ArrayList<>();
    	LeftIterator itLeft = new LeftIterator(root);
    	RightIterator itRight = new RightIterator(root);		
    	TreeNode cur = root;
    	TreeNode closestNode = root;
    	if (root == null) return res;
    	while (cur != null) {
    		if (cur.val == target) {
    			closestNode = cur;
    			break;
    		}
    		if (Math.abs(closestNode.val - target) > Math.abs(cur.val - target)) {
    			closestNode = cur;
    		}
    		if (cur.val < target) {
    			cur = cur.right;
    		} else {
    			cur = cur.left;
    		}
    	}
    	while (itLeft.hasNext()) {
    		if (itLeft.next() == closestNode.val) break;
    	}
    	while (itRight.hasNext()) {
    		if (itRight.next() == closestNode.val) break;
    	}
    	
    	res.add(closestNode.val);
    	for (int i = 1; i < k; i++) {
    		if (itLeft.hasNext() && itRight.hasNext()) {
    			if (target - itLeft.peek() < itRight.peek() - target) {
    				res.add(itLeft.next());
    			} else {
    				res.add(itRight.next());
    			}
    		} else if (itLeft.hasNext()) {
    			res.add(itLeft.next());
    		} else {
    			res.add(itRight.next());
    		}
    	}
    	return res;
    }
    static class RightIterator {
    	// is in-order iterator
    	Deque<TreeNode> stack;
    	public RightIterator(TreeNode root) {
    		stack = new LinkedList<>();
    		pushLeft(root);
    	}
    	private void pushLeft(TreeNode cur) {
    		while (cur != null) {
    			stack.push(cur);
    			cur = cur.left;
    		}
    	}		
    	public Integer next() {
    		if (!hasNext()) {
    			throw new IllegalStateException();
    		}
    		TreeNode cur = stack.pop();
    		pushLeft(cur.right);
    		return cur.val;
    	}
    	public boolean hasNext() {
    		return !stack.isEmpty();
    	}
    	public Integer peek() {
    		return stack.peek().val;
    	}
    }
    static class LeftIterator {
    	// reverse in-order iterator,  right, root, left
    	Deque<TreeNode> stack;
    	public LeftIterator(TreeNode root) {
    		stack = new LinkedList<>();
    		pushRight(root);
    	}
    	private void pushRight(TreeNode root) {
    		while (root != null) {
    			stack.push(root);
    			root = root.right;
    		}
    	}
    	public Integer next() {
    		if (!hasNext()) {
    			throw new IllegalStateException();
    		}
    		TreeNode cur = stack.pop();
    		pushRight(cur.left);
    		return cur.val;
    	}
    	public boolean hasNext() {
    		return !stack.isEmpty();
    	}
    	public Integer peek() {
    		return stack.peek().val;
    	}
    }
    

    ...


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.