flattenToDoublyLinkedListAsPreorder.java

public class flattenToDoublyLinkedListAsPreorder {
	
	private TreeNode head = null;
    private TreeNode tail = null;
    
    public TreeNode flattenRecursively(TreeNode root) {
        preOrder(root);
        return head;
    }
    
    private void preOrder(TreeNode node) {
        if (node == null)
            return;
        TreeNode leftCopy = node.left;
        TreeNode rightCopy = node.right;
        if(head == null) {
            head = node;
            head.left = null;
            tail = node;
            tail.right = null;
        } else {
            tail.right = node;
            node.left = tail;
            tail = tail.right;
            tail.right = null;
        }
        preOrder(leftCopy);
        preOrder(rightCopy);
    }
    
    public TreeNode flattenIteratively(TreeNode root) {
        TreeNode curr = root, head = null, tail = null;
        Stack<TreeNode> path = new Stack<TreeNode>();
        Stack<TreeNode> rightChildren = new Stack<TreeNode>();
        while (!path.isEmpty() || curr != null) {
            while (curr != null) {
                path.push(curr);
                rightChildren.push(curr.right);
                TreeNode leftCopy = curr.left;
                if (head == null) {
                    head = curr;
                    tail = curr;
                    head.left = null;
                    head.right = null;
                } else {
                    tail.right = curr;
                    curr.left = tail;
                    tail = tail.right;
                    tail.right = null;
                }
                curr = leftCopy;
            }
            path.pop();
            curr = rightChildren.pop();
        }
        return head;
    }
    
    //for test use, serialize doubly linked list
    public String serialize(TreeNode head) {
    	String res = "";
    	while(head != null) {
    		res += "<-";
    		res += head.val;
    		res += "->";
    		head = head.right;
    	}
    	return res;
    }
    
    public static void main(String[] args) {
    	test1();
    	test2();
    	test3();
    	test4();
    	test5();
    	test6();
    }
    
    //test cases
    private static void test1() {
    	TreeNode root = new TreeNode(1);
    	flattenToDoublyLinkedListAsPreorder f = new flattenToDoublyLinkedListAsPreorder();
    	if (f.serialize(f.flattenRecursively(root)).equals("<-1->"))
    		System.out.println("test case 1 passed!");
    	else
    		System.out.println("test case 1 failed!");
    }
    
    private static void test2() {
    	TreeNode root = new TreeNode(1);
    	flattenToDoublyLinkedListAsPreorder f = new flattenToDoublyLinkedListAsPreorder();
    	if (f.serialize(f.flattenIteratively(root)).equals("<-1->"))
    		System.out.println("test case 2 passed!");
    	else
    		System.out.println("test case 2 failed!");
    }
    
    private static void test3() {
    	TreeNode root = new TreeNode(1);
    	root.left = new TreeNode(2);
    	root.right = new TreeNode(3);
    	root.left.right = new TreeNode(4);
    	root.right.right = new TreeNode(5);
    	root.right.right.left = new TreeNode(6);
    	flattenToDoublyLinkedListAsPreorder f = new flattenToDoublyLinkedListAsPreorder();
    	if (f.serialize(f.flattenRecursively(root)).equals("<-1-><-2-><-4-><-3-><-5-><-6->"))
    		System.out.println("test case 3 passed!");
    	else
    		System.out.println("test case 3 failed!");
    }
    
    private static void test4() {
    	TreeNode root = new TreeNode(1);
    	root.left = new TreeNode(2);
    	root.right = new TreeNode(3);
    	root.left.right = new TreeNode(4);
    	root.right.right = new TreeNode(5);
    	root.right.right.left = new TreeNode(6);
    	flattenToDoublyLinkedListAsPreorder f = new flattenToDoublyLinkedListAsPreorder();
    	if (f.serialize(f.flattenIteratively(root)).equals("<-1-><-2-><-4-><-3-><-5-><-6->"))
    		System.out.println("test case 4 passed!");
    	else
    		System.out.println("test case 4 failed!");
    }
    
    private static void test5() {
    	TreeNode root = new TreeNode(3);
    	root.left = new TreeNode(7);
    	root.right = new TreeNode(6);
    	root.left.left = new TreeNode(4);
    	root.left.right = new TreeNode(2);
    	root.right.right = new TreeNode(5);
    	root.left.left.left = new TreeNode(12);
    	root.left.left.left.right = new TreeNode(0);
    	root.left.right.right = new TreeNode(10);
    	root.right.right.left = new TreeNode(8);
    	root.right.right.right = new TreeNode(9);
    	root.right.right.left.left = new TreeNode(17);
    	root.right.right.right.right= new TreeNode(20);
    	flattenToDoublyLinkedListAsPreorder f = new flattenToDoublyLinkedListAsPreorder();
    	if (f.serialize(f.flattenRecursively(root)).equals("<-3-><-7-><-4-><-12-><-0-><-2-><-10-><-6-><-5-><-8-><-17-><-9-><-20->"))
    		System.out.println("test case 5 passed!");
    	else
    		System.out.println("test case 5 failed!");
    }
    
    private static void test6() {
    	TreeNode root = new TreeNode(3);
    	root.left = new TreeNode(7);
    	root.right = new TreeNode(6);
    	root.left.left = new TreeNode(4);
    	root.left.right = new TreeNode(2);
    	root.right.right = new TreeNode(5);
    	root.left.left.left = new TreeNode(12);
    	root.left.left.left.right = new TreeNode(0);
    	root.left.right.right = new TreeNode(10);
    	root.right.right.left = new TreeNode(8);
    	root.right.right.right = new TreeNode(9);
    	root.right.right.left.left = new TreeNode(17);
    	root.right.right.right.right= new TreeNode(20);
    	flattenToDoublyLinkedListAsPreorder f = new flattenToDoublyLinkedListAsPreorder();
    	if (f.serialize(f.flattenIteratively(root)).equals("<-3-><-7-><-4-><-12-><-0-><-2-><-10-><-6-><-5-><-8-><-17-><-9-><-20->"))
    		System.out.println("test case 6 passed!");
    	else
    		System.out.println("test case 6 failed!");
    }
}