Skip to content

Instantly share code, notes, and snippets.

@hackwaly

hackwaly/treap.ml

Created May 26, 2023 22:22
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save hackwaly/4def289c5564605eea9f0d556fbf0205 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save hackwaly/4def289c5564605eea9f0d556fbf0205 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
treap.ml
type ('data, 'index) node =
{ left : ('data, 'index) node option
; right : ('data, 'index) node option
; priority : int
; mutable data : 'data
; mutable index : 'index
}
type ('a, 'b) t = ('a, 'b) node option
module type Build = sig
type data
type index
val build_index : data -> index option -> index option -> index
val append_data : data -> data -> data option
end
module type Dimension = sig
type index
val length_from_index : index -> int
end
module type Slice = sig
include Build
include Dimension with type index := index
val length : data -> int
val slice : data -> int -> int -> data
end
let get_left n = n.left
let get_right n = n.right
let get_data n = n.data
let get_index n = n.index
let empty = None
let fix
(type data index)
(module M : Build with type data = data and type index = index)
node
=
let index { index } = index in
let ( |?> ) o f = Option.map f o in
node.index <- M.build_index node.data (node.left |?> index) (node.right |?> index);
node
;;
let make
(type data index)
(module M : Build with type data = data and type index = index)
data
=
Some
{ left = None
; right = None
; priority = Random.full_int Int.max_int
; data
; index = M.build_index data None None
}
;;
let length (type index) (module M : Dimension with type index = index) treap =
match treap with
| None -> 0
| Some node -> M.length_from_index node.index
;;
type split_hack =
{ mutable f :
'data 'index.
(module Slice with type data = 'data and type index = 'index)
-> ('data, 'index) t
-> int
-> bool
-> ('data, 'index) t
}
type merge_hack =
{ mutable f :
'data 'index.
(module Build with type data = 'data and type index = 'index)
-> ('data, 'index) t
-> ('data, 'index) t
-> ('data, 'index) t
}
let split_hack : split_hack = { f = (fun _ -> assert false) }
let merge_hack : merge_hack = { f = (fun _ -> assert false) }
let split
(type a b)
(module M : Slice with type data = a and type index = b)
treap
index
return_left
=
match treap with
| Some n ->
let llen = length (module M) n.left in
if index <= llen
then split_hack.f (module M) n.left index return_left
else (
let i = index - llen in
let dlen = M.length n.data in
if i < dlen
then
if return_left
then merge_hack.f (module M) n.left (make (module M) (M.slice n.data 0 i))
else merge_hack.f (module M) (make (module M) (M.slice n.data i dlen)) n.right
else split_hack.f (module M) n.right (i - dlen) return_left)
| None ->
assert (index = 0);
None
;;
split_hack.f <- split
let sub
(type a b)
(module M : Slice with type data = a and type index = b)
treap
start
end_
=
let left = split (module M) treap end_ true in
split (module M) left start false
;;
let merge (type a b) (module M : Build with type data = a and type index = b) left right =
match left, right with
| Some ln, Some rn ->
Some
(if ln.priority <= rn.priority
then fix (module M) { ln with right = merge_hack.f (module M) ln.right right }
else fix (module M) { rn with left = merge_hack.f (module M) rn.left left })
| None, None -> None
| None, t | t, None -> t
;;
merge_hack.f <- merge
let append (type a b) (module M : Build with type data = a and type index = b) left right =
let rec step1 ln rn =
let rec aux n =
match n with
| { left = None; data; right } -> step2 ln data right
| { left = Some ln } -> aux ln
in
aux rn
and step2 ln data2 right2 =
let rec aux n =
match n with
| { right = None; data } ->
(match M.append_data data data2 with
| Some data ->
n.data <- data;
fix (module M) n |> ignore;
merge (module M) (Some n) right2
| None -> merge (module M) left right)
| { right = Some rn } -> aux rn
in
aux ln
in
match left, right with
| None, None -> None
| Some ln, Some rn -> step1 ln rn
| Some _, t | t, Some _ -> t
;;
Raw
treap.mli
type ('data, 'index) node
type ('a, 'b) t = ('a, 'b) node option
module type Build = sig
type data
type index
val build_index : data -> index option -> index option -> index
val append_data : data -> data -> data option
end
module type Dimension = sig
type index
val length_from_index : index -> int
end
module type Slice = sig
include Build
include Dimension with type index := index
val length : data -> int
val slice : data -> int -> int -> data
end
val get_left : ('a, 'b) node -> ('a, 'b) t
val get_right : ('a, 'b) node -> ('a, 'b) t
val get_data : ('a, 'b) node -> 'a
val get_index : ('a, 'b) node -> 'b
val empty : ('a, 'b) t
val make : 'a 'b. (module Build with type data = 'a and type index = 'b) -> 'a -> ('a, 'b) t
val length : 'a 'b. (module Dimension with type index = 'b) -> ('a, 'b) t -> int
val sub : 'a 'b. (module Slice with type data = 'a and type index = 'b) -> ('a, 'b) t -> int -> int -> ('a, 'b) t
val append : 'a 'b. (module Build with type data = 'a and type index = 'b) -> ('a, 'b) t -> ('a, 'b) t -> ('a, 'b) t
Raw
varray.ml
module type S = sig
type element
type t
val empty : t
val repeated : element -> int -> t
val listed : element array -> t
val length : t -> int
val append : t -> t -> t
val sub : t -> int -> int -> t
val get : t -> int -> element
end
module Make (Element : sig
type t
val equal : t -> t -> bool
end) : S with type element = Element.t = struct
type element = Element.t
type data =
| Repeated of element * int
| Listed of element array
type index = { length : int }
type t = (data, index) Treap.t
module M = struct
type nonrec data = data
type nonrec index = index
let length_from_index index = index.length
let build_index data left right =
let len =
match data with
| Repeated (_, c) -> c
| Listed a -> Array.length a
in
let llen = Option.(left |> map length_from_index |> value ~default:0) in
let rlen = Option.(right |> map length_from_index |> value ~default:0) in
{ length = len + llen + rlen }
;;
let append_data data1 data2 =
match data1, data2 with
| Repeated (e1, c1), Repeated (e2, c2) when Element.equal e1 e2 ->
Some (Repeated (e1, c1 + c2))
| Listed a1, Listed a2 when Array.length a1 + Array.length a2 <= 32 ->
Some (Listed (Array.append a1 a2))
| _ -> None
;;
let length data =
match data with
| Repeated (_, c) -> c
| Listed a -> Array.length a
;;
let slice data start end_ =
match data with
| Repeated (e, _) -> Repeated (e, end_ - start)
| Listed a -> Listed (Array.sub a start end_)
;;
end
let empty = Treap.empty
let repeated elt count = Treap.make (module M) (Repeated (elt, count))
let listed arr = Treap.make (module M) (Listed arr)
let length va = Treap.length (module M) va
let append va1 va2 = Treap.append (module M) va1 va2
let sub va start end_ = Treap.sub (module M) va start end_
let get va index =
let rec aux t index =
match t with
| Some n ->
let lt = Treap.get_left n in
let llen = length lt in
if index < llen
then aux lt index
else (
let index = index - llen in
let data = Treap.get_data n in
let dlen = M.length data in
if index < dlen
then (
match data with
| Repeated (e, _) -> e
| Listed a -> Array.get a index)
else (
let rt = Treap.get_right n in
aux rt (index - dlen)))
| None -> assert false
in
aux va index
;;
end
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment