SHARE
TWEET

Untitled

a guest Nov 21st, 2019 89 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. type t = Empty | Node of t * (int * int) * t * (int * int) * int
  2. exception Impossible
  3. open Printf
  4.  
  5. let empty = Empty
  6. let is_empty avl = avl = Empty
  7.  
  8. let numbers_between (a, b) = b - a + 1
  9. let numbers_of x = match x with | Empty -> (0, 0) | Node (_, _, _, a, _) -> a
  10. let min_of a = match a with | Empty -> min_int | Node (_, (b, c), _, _, _) -> b
  11. let height_of x = match x with | Empty -> 0 | Node (_, _, _, _, a) -> a
  12. let make_node l v r = Node (l, v, r,
  13.   (let (ln, li) = numbers_of l in
  14.   let (rn, ri) = numbers_of r in
  15.     if numbers_between v = max_int then (0, 1)
  16.     else if numbers_between v <= 0 then (0, 1)
  17.     else if li = 2 || ri = 2 then (min_int, 2)
  18.     else if li = 1 || ri = 1 then (ln + rn + numbers_between v, 1)
  19.     else if min_int + ln + rn + numbers_between v = max_int then (0, 1)
  20.     else if min_int + ln + rn + numbers_between v >= 0 then (min_int + ln + rn + numbers_between v, 1)
  21.     else if min_int + ln + rn + numbers_between v = min_int then (min_int + ln + rn + numbers_between v, 2)
  22.     else (ln + rn + numbers_between v, 0)),
  23.   (max (height_of l) (height_of r)) + 1)
  24.  
  25. let rec fold f avl acc = match avl with | Empty -> acc | Node (l, (a, b), r, _, _) -> fold f r (f (a, b) (fold f l acc))
  26. let rec iter f avl = match avl with | Empty -> () | Node (l, v, r, _, _) -> iter f l; f v; iter f r
  27.  
  28. let bal l k r =
  29.   let hl = height_of l in
  30.   let hr = height_of r in
  31.   if hl > hr + 2 then
  32.     match l with
  33.     | Node (ll, lk, lr, _, _) ->
  34.         if height_of ll >= height_of lr then make_node ll lk (make_node lr k r)
  35.         else
  36.           (match lr with
  37.           | Node (lrl, lrk, lrr, _, _) ->
  38.               make_node (make_node ll lk lrl) lrk (make_node lrr k r)
  39.           | Empty -> assert false)
  40.     | Empty -> assert false
  41.   else if hr > hl + 2 then
  42.     match r with
  43.     | Node (rl, rk, rr, _, _) ->
  44.         if height_of rr >= height_of rl then make_node (make_node l k rl) rk rr
  45.         else
  46.           (match rl with
  47.           | Node (rll, rlk, rlr, _, _) ->
  48.               make_node (make_node l k rll) rlk (make_node rlr rk rr)
  49.           | Empty -> assert false)
  50.     | Empty -> assert false
  51.   else make_node l k r
  52.  
  53. let cmp (a, b) (c, d) =
  54.   if b = c - 1 && c <> min_int then -3 else
  55.   if b < c then -2 else
  56.   if b = c then -1 else
  57.   if b > c && b < d then 0 else
  58.   if b = d then 1 else
  59.   if b = d + 1 && d <> max_int then 3 else
  60.   if b > d then 2 else 4
  61.  
  62. let rec add_one x t =
  63. match t with
  64. | Node (l, k, r, els, h) ->
  65.     let c = cmp x k in
  66.     if c < 0 then
  67.       let nl = add_one x l in
  68.       bal nl k r
  69.     else if c > 0 then
  70.       let nr = add_one x r in
  71.       bal l k nr
  72.     else raise Impossible
  73. | Empty -> make_node Empty x Empty
  74.  
  75. let rec join l v r =
  76.   match (l, r) with
  77.   | (Empty, _) -> add_one v r
  78.   | (_, Empty) -> add_one v l
  79.   | (Node(ll, lv, lr, el_l, lh), Node(rl, rv, rr, el_r, rh)) ->
  80.       if lh > rh + 2 then bal ll lv (join lr v r) else
  81.       if rh > lh + 2 then bal (join l v rl) rv rr else
  82.       make_node l v r
  83.  
  84.  
  85.   let rec merge sm gr =
  86.     match sm with
  87.     | Node (l, v, r, n, h) -> let nr = (merge r gr) in bal l v nr
  88.     | Empty -> gr
  89.  
  90. let split_help x t =
  91.   let clear = (Empty, false, Empty, x, x) in
  92.     let rec loop x t (s, c, g, l_a, r_a) =
  93.       match t with
  94.       |  Empty -> (s, c, g, l_a, r_a)
  95.       | Node (l, (a, b), r, _, _) ->
  96.           let d = cmp (x, x) (a, b) in
  97.           if d = -3 then
  98.             let (ss, cc, gg, l_aa, r_aa) = loop x l clear
  99.             in (merge s ss, false, merge g (merge gg r), l_aa, max (max r_a r_aa) b) else
  100.           if d = 3 then
  101.             let (ss, cc, gg, l_aa, r_aa) = loop x r clear
  102.             in (merge s (merge l ss), false, merge gg g, min (min l_a l_aa) a, r_aa) else
  103.           if d = -2 then
  104.             let (ss, cc, gg, l_aa, r_aa) = loop x l clear
  105.             in (merge s ss, cc || c, merge g (join gg (a, b) r), l_aa, r_aa) else
  106.           if d = 2 then
  107.             let (ss, cc, gg, l_aa, r_aa) = loop x r clear
  108.             in (merge s (join l (a, b) ss), cc || c, merge gg g, l_aa, r_aa) else
  109.           if d = -1 then
  110.             (merge s l, true, merge r g, x, b) else
  111.           if d = 1 then
  112.             (merge s l, true, merge r g, a, x) else
  113.           if d = 0 then
  114.             (merge s l, true, merge r g, a, b) else raise Impossible
  115.   in loop x t clear
  116.  
  117. let add (x, y) t =
  118.   let (s, _, _, a, _) = split_help x t in
  119.   let (_, _, g, _, b) = split_help y t in
  120.   add_one (a, b) (merge s g)
  121.  
  122. let remove (x, y) t =
  123.   let (s, _, _, a, _) = split_help x t in
  124.   let (_, _, g, _, b) = split_help y t in
  125.   let s = if x - a > 0 then add_one (a, x - 1) s else s in
  126.   let g = if b - y > 0 then add_one (y + 1, b) g else g in
  127.   merge s g
  128.  
  129. let split x t =
  130.   let (s, c, g, l_a, r_a) = split_help x t in
  131.   let g = (if r_a > x then (add_one (x + 1, r_a) g) else g) in
  132.   let s = (if l_a < x then (add_one (l_a, x - 1) s) else s) in
  133.   (s, c, g)
  134.  
  135. let elements x =
  136.   let rec loop acc x =
  137.     match x with
  138.     | Empty -> acc
  139.     | Node(l, k, r, _, _) -> loop (k :: (loop acc r)) l in
  140.   loop [] x
  141.  
  142. let rec mem x t =
  143.   match t with
  144.   | Empty -> false
  145.   | Node (l, (a, b), r, _, _) ->
  146.     if x >= a && x <= b then true
  147.     else if x < a then mem x l else mem x r
  148.  
  149. let rec below_help x t s =
  150.   match t with
  151.   | Empty -> (0, 0)
  152.   | Node (l, (a, b), r, (n, i), h) ->
  153.     if x > b then
  154.       (let (ln, li) = numbers_of l in
  155.       let (b_n, b_i) = below_help x r 0 in
  156.       if b_i > 0 then (max_int, 1)
  157.       else if li > 0 then (max_int, 1)
  158.       else if min_int + b_n + ln + b - a + 1 >= 0 then (max_int, 1)
  159.       else if min_int + b_n + ln + b - a + 1 = min_int then (max_int, 2)
  160.       else (b_n + ln + b - a + 1, 0))
  161.     else if x >= a then
  162.       (let (ln, li) = numbers_of l in
  163.       if li > 0 then (max_int, 1)
  164.       else if min_int + ln + x - a + 1 >= 0 then (max_int, 1)
  165.       else if min_int + ln + x - a + 1 = min_int then (max_int, 2)
  166.       else (ln + x - a + 1, 0))
  167.     else
  168.       (let (aa, ii) = below_help x l s in
  169.       if ii > 0 then (max_int, 1)
  170.       else if min_int + aa >= 0 then (max_int, 1) else (aa, 0))
  171.  
  172. let below x t =
  173.   let (ans, _) = below_help x t 0 in ans
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
Not a member of Pastebin yet?
Sign Up, it unlocks many cool features!
 
Top