AVL - ith element, C++ (gcc)

Language: Layout:

#include <bits/stdc++.h>

using namespace std;
/*

1 2 3 4 5 6

1 
 2
  3
   4 
    5
     6
     
     3
    2 5
   1   6
   
      _
   x     x
   [abs(L_SZ - R_SZ) <= 2] --> Almost equal
   
        N
    N/2   N/2
   N/4   N/4
  N/8  N/8
 N/16 ..

N/2^x = 1
     
O(N)

O(logN)

[T] X [T+2]
     Y     X1
   Z   Y1   \\\\

Y
     Z   X
        y1 X1

[T+2] X [T]
      Y.  Z
         Z1 Z2

Z
      X    Z2
    Y  Z1

Balance factor = (lft_size - rgt_size)

[T+2]   X  [T] -> [2]
        
 [1] Y      Z  -> [-1]
 
  Y1 Y2.  Z1 Z2

Y [0]
  [T]Y1       X [0] [T]
        [T-1]Y2  Z (-1) [T-1]
             ....

=============================================

[T+2]   X  [T] -> [2]
        
 [-1]Y      Z  -> [1]
 
  Y1 Y2.  Z1 Z2

Y [-2]
  [T-1]Y1       X [1] [T+1]
          [T]Y2  Z (+1) [T-1]
             ....
             
            
            
1. Balance factor
2. Not almost eq node
3. Decide where to rotate (depend which is heavier)
4. Compre the sign of the node which is going to be new root with the current root's balance factor
    - Make the sign same with appropraite rotation
5. Do the rotation

struct node {
    node *lft, *rgt;
    int val;
    
    int bal;
    int hei;
    int siz;
    
    node(int v) : val(v), lft(NULL), rgt(NULL), bal(0), hei(1), siz(1) {}
};

int height (node* root) {
    if (!root) return 0;
    
    int lft = root -> lft ? root -> lft -> hei : 0;
    int rgt = root -> rgt ? root -> rgt -> hei : 0;
    
    return (max(lft, rgt) + 1);
}

node* compute (node* root) {
    root -> hei = height(root);
    
    int lft_hei = height(root -> lft);
    int rgt_hei = height(root -> rgt);
    root -> bal = lft_hei - rgt_hei;
    
    int lft_sz = root -> lft? root -> lft -> siz: 0;
    int rgt_sz = root -> rgt? root -> rgt -> siz: 0;
    root -> siz = lft_sz + rgt_sz + 1;
    
    return root;
}

/*
  [T+2] X [T]
      Y.  Z
         Z1 Z2

Z
      X    Z2
    Y  Z1
*/
node* lft_rotate(node* root) {
    node* new_root = root -> rgt;
    
    root -> rgt = new_root -> lft;
    new_root -> lft = root;
    
    compute (root);
    compute (new_root);
    
    return new_root;
}

[T] X [T+2]
     Y     X1
   Z   Y1   \\\\

Y
     Z   X
        y1 X1

*/
node* rgt_rotate(node* root) {
    node* new_root = root -> lft;
    
    root -> lft = new_root -> rgt;
    new_root -> rgt = root;
    
    compute (root);
    compute (new_root);
    
    return new_root;
}

/*
4. Compre the sign of the node which is going to be new root with the current root's balance factor
    - Make the sign same with appropraite rotation
5. Do the rotation
*/
node* balance(node* root) {
    if (abs(root -> bal) <= 1) return root;
    
    if (root -> bal == 2) {
        if (root -> lft -> bal == -1) root -> lft = lft_rotate(root -> lft);
        return rgt_rotate(root);
    }
    else {
        // -2
        if (root -> rgt -> bal == 1) root -> rgt = rgt_rotate(root -> rgt);
        return lft_rotate(root);
    }
}

node* insert(node* root, int new_val) {
    if (!root) {
        return new node(new_val);
    }
    
    if (root -> val < new_val) {
        root -> rgt = insert (root -> rgt, new_val);
    }
    else {
        root -> lft = insert (root -> lft, new_val);
    }
    
    compute (root);
    
    return balance(root);
}

int query (node* root, int pos) {
    int lft_sz = root -> lft? root -> lft -> siz: 0;
    int rgt_sz = root -> rgt? root -> rgt -> siz: 0;
    
    if (lft_sz+1 == pos) return root -> val;
    if (lft_sz+1 < pos) {
        return query (root -> rgt, pos-lft_sz-1);
    }
    else {
        return query (root -> lft, pos);
    }
}

int main () {
    node* avl = NULL;
    
    for (int j = 1; j <= 100; j ++) avl = insert(avl, j*2);
    for (int j = 1; j < 10; j ++) {
        int pos = rand() % 100 + 1;
        cout << pos << " --> " << query(avl, pos) << endl;
    }
    return 0;
}

Show compiler warnings [ + ] Compiler args [ + ] Show input
Compilation time: 2.55 sec, absolute running time: 0.27 sec, cpu time: 0.02 sec, memory peak: 3 Mb, absolute service time: 2,88 sec