紅黑樹的簡易實現

肝了好幾個小時的成品
大概經過了洛谷11/12個測試點(其中一個TLE,時限開的太緊了)

簡要說明:
1.add的旋轉分類參考自算法第4版實現(大概就是右red/左兩個red/一左一右red)
2.刪除使用了LAZY標記,分類刪除?不存在的

題面

您須要寫一種數據結構（可參考題目標題），來維護一些數，其中須要提供如下操做：

插入xxx數
刪除xxx數(如有多個相同的數，因只刪除一個)
查詢xxx數的排名(排名定義爲比當前數小的數的個數+1+1+1。如有多個相同的數，因輸出最小的排名)
查詢排名爲xxx的數
求xxx的前驅(前驅定義爲小於xxx，且最大的數)
求xxx的後繼(後繼定義爲大於xxx，且最小的數)

輸入格式

第一行爲nnn，表示操做的個數,下面nnn行每行有兩個數optoptopt和xxx，optoptopt表示操做的序號( 1≤opt≤6 1 \leq opt \leq 6 1≤opt≤6 )
輸出格式

對於操做3,4,5,63,4,5,63,4,5,6每行輸出一個數，表示對應答案

輸入樣例

10
1 106465
4 1
1 317721
1 460929
1 644985
1 84185
1 89851
6 81968
1 492737
5 493598

輸出樣例

106465
84185
492737

import java.io.*;
import java.util.ArrayList;
import java.util.List;
import java.util.StringTokenizer;

enum Color {
    RED,
    BLACK;
}

class Node implements Comparable<Node> {

    public int key;
    public int size;
    public int cnt;
    public Color color;
    public Node lc,rc;

    public Node(int key,int size,Color color) {
        this.key = key;
        this.size = size;
        this.color = color;
        this.cnt = 1;
    }

    /** 能夠把int key改成Comparable **/
    public int compareTo(Node that) {
        return this.key - that.key;
    }

    public String toString() {
        return ""+key;
    }

}
class RedBlackTree {

    public Node root;

    private static Color RED = Color.RED;
    private static Color BLACK = Color.BLACK;

    /** rotate right **/
    private Node rr(Node n) {
        Node fa = n.lc;
        n.lc = fa.rc;
        fa.rc = n;

        fa.color = n.color;
        n.color = RED;

        pushUp(n);
        pushUp(fa);
        return fa;
    }

    /** rotate left **/
    private Node rl(Node n) {
        Node fa = n.rc;
        n.rc = fa.lc;
        fa.lc = n;

        fa.color = n.color;
        n.color = RED;

        pushUp(n);
        pushUp(fa);
        return fa;
    }

    private Color flippedColor(Node n) {
        return isRed(n) ? BLACK : RED;
    }

    private void flip(Node n) {
        n.color = flippedColor(n);
        n.lc.color = flippedColor(n.lc);
        n.rc.color = flippedColor(n.rc);
    }

    private boolean isRed(Node n) {
        return n != null && n.color == RED;
    }

    private int size(Node n) {
        return n == null ? 0 : n.size;
    }

    private void pushUp(Node n) {
        if(n == null) return;
        n.size = n.cnt + size(n.lc) + size(n.rc);
    }

    public void add(int key) {
        if(root == null) {
            root = new Node(key,1,RED);
            return;
        }
        root = add(root,key);
        root.color = BLACK;
    }
    private Node add(Node n,int key) {
        if(n == null) return new Node(key,1,RED);
        int cmp = key-n.key;
        if(cmp < 0) n.lc = add(n.lc,key);
        else if(cmp > 0) n.rc = add(n.rc,key);
        else {
            n.cnt++;
            n.size++;
            pushUp(n);
            return n;
        }

        if(isRed(n.rc) && !isRed(n.lc)) n = rl(n);
        if(isRed(n.lc) && isRed(n.lc.lc)) n = rr(n);
        if(isRed(n.lc) && isRed(n.rc)) flip(n);

        pushUp(n);
        return n;
    }

    public boolean get(int key) {
        Node cur = root;
        while(cur != null) {
            int cmp = key-cur.key;
            if(cmp < 0) cur = cur.lc;
            else if(cmp > 0) cur = cur.rc;
            else return true;
        }
        return false;
    }

    public boolean remove(int key) {
        List<Node> stack = new ArrayList<>();
        Node cur = root;
        while(cur != null) {
            int cmp = key-cur.key;
            stack.add(cur);
            if(cmp < 0) cur = cur.lc;
            else if(cmp > 0) cur = cur.rc;
            else {
                if(cur.cnt == 0) return false;
                cur.size--;
                cur.cnt--;
                for(int i = stack.size()-1; i >= 0; --i) {
                    pushUp(stack.get(i));
                }
                return true;
            }
        }
        return false;
    }

    public Node findNode(int key) {
        Node cur = root;
        while(cur != null) {

            int cmp = key-cur.key;
            if(cmp < 0) cur = cur.lc;
            else if(cmp > 0) cur = cur.rc;
            else break;
        }
        return cur;
    }

    public int getRank(int key) {
        Node cur = root;
        int tmp = 0;
        while(cur != null) {
            int cmp = key-cur.key;
            if(cmp < 0) {
                cur = cur.lc;
            } else if(cmp > 0) {
                tmp += cur.cnt + size(cur.lc);
                cur = cur.rc;
            } else {
                return size(cur.lc)+1 + tmp;
            }
        }
        return 0;
    }

    public Node getKth(int kth) {
        Node cur = root;
        while(cur != null) {

            int sz = size(cur.lc);
            if(sz >= kth) {
                cur = cur.lc;
            } else if(kth-sz > cur.cnt) {
                kth -= (sz + cur.cnt);
                cur = cur.rc;
            } else {
                return cur;
            }
        }
        return null;
    }

    /** 小於key且最大的數 **/
    public int lower(int key) {
        return lower(root,key);
    }

    public int lower(Node cur,int key) {
        int result = Integer.MIN_VALUE;
        if(cur == null) return result;
        while(cur != null) {
            int cmp = key-cur.key;
            if(cmp > 0 && cur.key > result) {
                if(cur.cnt > 0) result = cur.key;
                else return Math.max(result,Math.max(lower(cur.lc,key),lower(cur.rc,key)));
            } else if(cmp > 0) {
                cur = cur.rc;
            } else {
                cur = cur.lc;
            }
        }
        return result;
    }

    /** 大於key且最小的數 **/
    public int upper(int key) {
        return upper(root,key);
    }

    /** 被LAZY刪除的數可能會致使死循環,須要分支遍歷 **/
    public int upper(Node cur,int key) {
        int result = Integer.MAX_VALUE;
        while(cur != null) {
            int cmp = key-cur.key;
            if(cmp < 0 && cur.key < result) {
                if(cur.cnt > 0) result = cur.key;
                else return Math.min(result,Math.min(upper(cur.lc,key),upper(cur.rc,key))); // 注意
            } else if(cmp < 0) {
                cur = cur.lc;
            } else {
                cur = cur.rc;
            }
        }
        return result;
    }



    public void dfs(Node n) {
        if(n == null) return;
        dfs(n.lc);
        if(n.cnt > 0) System.out.print(n.key+" ");
        dfs(n.rc);
    }

    public void dfs2(Node n) {
        if(n == null) return;
        if(n.cnt > 0) System.out.print(n.key+"[size"+n.size+"] ");
        dfs2(n.lc);
        dfs2(n.rc);
    }


}

public class Main {

    /** 洛谷的時限卡的比較嚴,須要開個小掛 **/
    static class InputReader {
        public BufferedReader reader;
        public StringTokenizer tokenizer;

        public InputReader(InputStream stream) {
            reader = new BufferedReader(new InputStreamReader(stream), 32768);
            tokenizer = null;
        }

        public String next() {
            while (tokenizer == null || !tokenizer.hasMoreTokens()) {
                try {
                    tokenizer = new StringTokenizer(reader.readLine());
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }
            return tokenizer.nextToken();
        }

        public int nextInt() {
            return Integer.parseInt(next());
        }

    }



    public static void main(String[] args)  {
        RedBlackTree rbt = new RedBlackTree();
        InputReader sc = new InputReader(System.in);
        int n = sc.nextInt();
        for(int i = 1; i <= n; i++) {
            int op = sc.nextInt();
            int key = sc.nextInt();
            if(op == 1) {
                rbt.add(key);
            } else if(op == 2) {
                rbt.remove(key);
            } else if(op == 3) {
                System.out.println(rbt.getRank(key));
            } else if(op == 4) {
                System.out.println(rbt.getKth(key));
            } else if(op == 5) {
                System.out.println(rbt.lower(key));
            } else {
                System.out.println(rbt.upper(key));
            }
        }
    }
}