package com.company;import java.util.*;

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

class StkNode {
    TreeNode tn;
    int count;
    StkNode(TreeNode tn) {
        this.tn = tn;
        count = 0;
    }
}

class Solution {

    List<Integer> getArray(TreeNode root) {

        List<Integer> ret = new ArrayList<>();
        if (root.left != null) {
            ret.addAll(getArray(root.left));
        }
        ret.add(root.val);
        if (root.right != null) {
            ret.addAll(getArray(root.right));
        }
        return ret;
    }

    public void recoverTree(TreeNode root) {

        // Space O(N)的解法,想到了是直接将BST放平,然后重新排序,然后比较一下就可以。

        // 分情况, 左ok, 右ok, 1. 上ok, 2. 上不ok
        // 左不ok, 或者右不ok, 1. 上ok,
        // 想不清楚,先用直接的方法:
        if (root == null) {
            return;
        }
        List<Integer> ret = getArray(root);
        Integer[] sortRet= ret.toArray(new Integer[0]);
        Arrays.sort(sortRet);

        int[] change = new int[2];
        int[] index = new int[2];
        int ii = 0;
        for (int i=0; i<sortRet.length; i++) {
            //System.out.printf("old: %d, new: %d\n", ret.get(i), sortRet[i]);
            if (ret.get(i) != sortRet[i]) {
                index[ii] = i+1;
                change[ii] = sortRet[i];
                ii++;
                if (ii >= 2) {
                    break;
                }
            }
        }
        //System.out.printf("ii:%d\n", ii);

        Stack<StkNode> stk = new Stack<>();
        int count = 0;
        int k = 0;
        StkNode stkRoot = new StkNode(root);
        stk.push(stkRoot);
        while (!stk.isEmpty()) {
            StkNode tmp = stk.pop();
            //System.out.printf("here: %d, %d, %d\n", tmp.count, count, tmp.tn.val);
            if (tmp.count == 1) {
                count++;
                if (count == index[k]) {
                    //System.out.printf("here: %d\n", count);
                    tmp.tn.val = change[k];
                    k++;
                    if (k>=2) {
                        return;
                    }
                }
                if (tmp.tn.right != null) {
                    StkNode newTmp = new StkNode(tmp.tn.right);
                    stk.push(newTmp);
                }
            }
            else {
                tmp.count++;
                stk.push(tmp);
                if (tmp.tn.left != null) {
                    StkNode newTmp = new StkNode(tmp.tn.left);
                    stk.push(newTmp);
                }
            }
        }

    }
}

public class Main {

    public static void main(String[] args) {
        System.out.println("Hello!");
        Solution solution = new Solution();

        TreeNode root = new TreeNode(2);
        TreeNode root2 = new TreeNode(3);
        TreeNode root3 = new TreeNode(1);
        root.left = root2;
        root.right = root3;
        solution.recoverTree(root);
        System.out.printf("Get ret: \n");
        System.out.printf("Get ret1: %d\n", root.left.val);
        System.out.printf("Get ret2: %d\n", root.val);
        System.out.printf("Get ret3: %d\n", root.right.val);
        System.out.println();

        /*Iterator<List<Integer>> iterator = ret.iterator();
        while (iterator.hasNext()) {
            Iterator iter = iterator.next().iterator();
            while (iter.hasNext()) {
                System.out.printf("%d,", iter.next());
            }
            System.out.println();
        }*/

        System.out.println();

    }
}

https://leetcode.com/problems/recover-binary-search-tree/

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
    stack < pair <TreeNode*, int> > stk;
public:
    void recoverTree(TreeNode* root) {
        if (root == NULL) {
            return;
        }
        pair<TreeNode*, int> pr(root, 0);
        stk.push(pr);
        
        TreeNode * first_result = NULL;
        TreeNode * last_result = NULL;
        TreeNode* last = NULL;
        
        pair<TreeNode*, int> tmp;
        TreeNode* tn_tmp;

        while ( !stk.empty() ) {
            tmp = stk.top();
            stk.pop();
            
            if (tmp.second == 0) {
                // left
                tn_tmp = tmp.first;
                pair<TreeNode*, int> tmp_cur(tn_tmp, 1);
                stk.push(tmp_cur);
                    
                if (tn_tmp->left != NULL) {
                    pair<TreeNode*, int> tmp_left(tn_tmp->left, 0);
                    stk.push(tmp_left);
                }
                
            }
            else {
                // right
                tn_tmp = tmp.first;
                if (last != NULL && last->val > tn_tmp->val) {
                    // Found
                    if (first_result == NULL ) {
                        first_result = last;
                        last_result = tn_tmp;
                    }
                    else {
                        last_result = tn_tmp;
                        break;
                    }
                }
                last = tn_tmp;
                
                if (tn_tmp->right != NULL) {
                    pair<TreeNode*, int> tmp_pr(tn_tmp->right, 0);
                    stk.push(tmp_pr);
                }
                
            }
            
        }
        
        if (first_result != NULL && last_result != NULL) {
            int swap = first_result->val;
            first_result->val = last_result->val;
            last_result->val = swap;
        }
        return;
                
    }

};