1+ namespace Algorithms.DataStructures
2+
3+ module Treap =
4+ type TreapNode = {
5+ Value: int
6+ Priority: int64
7+ Size: int
8+ LeftChild: Option < TreapNode >
9+ RightChild: Option < TreapNode >
10+ }
11+
12+ module TreapNode =
13+ let create ( value : int ) : TreapNode =
14+ {
15+ Value = value
16+ Priority = System.Random() .NextInt64()
17+ Size = 1
18+ LeftChild = None
19+ RightChild = None
20+ }
21+
22+ let getSize ( maybeNode : Option < TreapNode >) : int =
23+ maybeNode
24+ |> Option.map ( fun node -> node.Size)
25+ |> Option.defaultValue 0
26+
27+ type TreapNode
28+ with
29+ member this.UpdateLeftChild ( leftChild : Option < TreapNode >) : TreapNode =
30+ {
31+ this with
32+ LeftChild = leftChild
33+ Size = 1 + TreapNode.getSize leftChild + TreapNode.getSize this.RightChild
34+ }
35+ member this.UpdateRightChild ( rightChild : Option < TreapNode >) : TreapNode =
36+ {
37+ this with
38+ RightChild = rightChild
39+ Size = 1 + TreapNode.getSize this.LeftChild + TreapNode.getSize rightChild
40+ }
41+
42+ [<RequireQualifiedAccess>]
43+ type Treap = {
44+ Root: Option < TreapNode >
45+ }
46+
47+ let empty : Treap = { Root = None }
48+
49+ /// Splits treap into two parts based on value
50+ /// Returns (left, right) tuple where:
51+ /// - left contains all nodes with values < split value
52+ /// - right contains all nodes with values >= split value
53+ let rec split ( root : Option < TreapNode >) ( value : int ) : Option < TreapNode > * Option < TreapNode > =
54+ match root with
55+ | None ->
56+ None, None
57+ | Some root ->
58+ if root.Value < value then
59+ // root node belongs to the left
60+ // split the right child of root and update the right child of root
61+ let updatedRightOfRoot , right = split root.RightChild value
62+ root.UpdateRightChild updatedRightOfRoot |> Some, right
63+ else
64+ // root node belongs to the right
65+ // split the left child of root and update the left child of root
66+ let left , updatedLeftOfRoot = split root.LeftChild value
67+ left, root.UpdateLeftChild updatedLeftOfRoot |> Some
68+
69+ /// Merges two treaps maintaining BST and heap properties
70+ /// Assumes all values in left treap are less than all values in right treap
71+ let rec merge ( left : Option < TreapNode >) ( right : Option < TreapNode >) : Option < TreapNode > =
72+ match left, right with
73+ | None, right -> right
74+ | left, None -> left
75+ | Some left, Some right ->
76+ if left.Priority < right.Priority then
77+ // left node is the root of the merged treap, merge its right child with right treap
78+ let updatedLeftsRightChild = merge left.RightChild ( Some right)
79+ left.UpdateRightChild updatedLeftsRightChild |> Some
80+ else
81+ // right node is the root of the merged treap, merge its left child with left treap
82+ let updatedRightsLeftChild = merge ( Some left) right.LeftChild
83+ right.UpdateLeftChild updatedRightsLeftChild |> Some
84+
85+ // Inserts a new value into the treap, noop if value already exists
86+ let insert ( value : int ) ( treap : Treap ) : Treap =
87+ let node = TreapNode.create value
88+ let left , right = split treap.Root value
89+ let _ , right = split right ( value + 1 )
90+ merge ( merge left ( Some node)) right
91+ |> fun root -> { Root = root }
92+
93+ let erase ( value : int ) ( treap : Treap ) : Treap =
94+ let left , right = split treap.Root value
95+ let _ , right = split right ( value + 1 )
96+ merge left right
97+ |> fun root -> { Root = root }
98+
99+ /// Gets the kth smallest element in the treap (0-indexed)
100+ /// Returns None if k is out of bounds
101+ let getKthElement ( treap : Treap ) ( k : uint ) : Option < int > =
102+ if TreapNode.getSize treap.Root |> uint <= k then
103+ None
104+ else
105+ let rec getKthElementImpl ( root : TreapNode ) ( k : int ) : int =
106+ assert ( k < root.Size)
107+ if root.Size = 1 then
108+ root.Value
109+ else
110+ if k < TreapNode.getSize root.LeftChild then
111+ getKthElementImpl ( root.LeftChild |> Option.get) k
112+ elif k = TreapNode.getSize root.LeftChild then
113+ root.Value
114+ else
115+ getKthElementImpl ( root.RightChild |> Option.get) ( k - TreapNode.getSize root.LeftChild - 1 )
116+ getKthElementImpl ( treap.Root |> Option.get) ( int k) |> Some
117+
118+ /// Gets the index of a value in the treap (0 indexed position in sorted order)
119+ /// Returns None if value is not found
120+ let getIndex ( treap : Treap ) ( value : int ) : Option < int > =
121+ let left , right = split treap.Root value
122+ let node , _right = split right ( value + 1 )
123+ node
124+ |> Option.map ( fun _ -> TreapNode.getSize left)
0 commit comments