@@ -25,146 +25,145 @@ THE SOFTWARE.*/
2525import Foundation
2626
2727public indirect enum Treap < Key: Comparable , Element> {
28- case Empty
29- case Node( key: Key , val: Element , p: Int , left: Treap , right: Treap )
30-
31- public init ( ) {
32- self = . Empty
33- }
34-
35- internal func get( key: Key ) -> Element ? {
36- switch self {
37- case . Empty:
38- return nil
39- case let . Node( treeKey, val, _, _, _) where treeKey == key:
40- return val
41- case let . Node( treeKey, _, _, left, _) where key < treeKey:
42- return left. get ( key)
43- case let . Node( treeKey, _, _, _, right) where key > treeKey:
44- return right. get ( key)
45- default :
46- return nil
47- }
28+ case empty
29+ case node( key: Key , val: Element , p: Int , left: Treap , right: Treap )
30+
31+ public init ( ) {
32+ self = . empty
33+ }
34+
35+ internal func get( _ key: Key ) -> Element ? {
36+ switch self {
37+ case . empty:
38+ return nil
39+ case let . node( treeKey, val, _, _, _) where treeKey == key:
40+ return val
41+ case let . node( treeKey, _, _, left, _) where key < treeKey:
42+ return left. get ( key)
43+ case let . node( treeKey, _, _, _, right) where key > treeKey:
44+ return right. get ( key)
45+ default :
46+ return nil
4847 }
49-
50- public func contains ( key : Key ) -> Bool {
51- switch self {
52- case . Empty :
53- return false
54- case let . Node ( treeKey , _ , _ , _ , _ ) where treeKey == key :
55- return true
56- case let . Node ( treeKey , _ , _ , left , _ ) where key < treeKey :
57- return left . contains ( key)
58- case let . Node ( treeKey , _ , _ , _ , right ) where key > treeKey :
59- return right . contains ( key)
60- default :
61- return false
62- }
48+ }
49+
50+ public func contains ( _ key : Key ) -> Bool {
51+ switch self {
52+ case . empty :
53+ return false
54+ case let . node ( treeKey , _ , _ , _ , _ ) where treeKey == key :
55+ return true
56+ case let . node ( treeKey , _ , _ , left , _ ) where key < treeKey :
57+ return left . contains ( key)
58+ case let . node ( treeKey , _ , _ , _ , right ) where key > treeKey :
59+ return right . contains ( key )
60+ default :
61+ return false
6362 }
64-
65- public var depth : Int {
66- get {
67- switch self {
68- case . Empty :
69- return 0
70- case let . Node ( _ , _ , _ , left , . Empty ) :
71- return 1 + left. depth
72- case let . Node ( _ , _ , _ , . Empty , right ) :
73- return 1 + right. depth
74- case let . Node ( _ , _ , _ , left , right) :
75- let leftDepth = left. depth
76- let rightDepth = right . depth
77- return 1 + max ( leftDepth , rightDepth )
78- }
79-
80- }
63+ }
64+
65+ public var depth : Int {
66+ get {
67+ switch self {
68+ case . empty :
69+ return 0
70+ case let . node ( _ , _ , _ , left, . empty ) :
71+ return 1 + left . depth
72+ case let . node ( _ , _ , _ , . empty , right) :
73+ return 1 + right. depth
74+ case let . node ( _ , _ , _ , left, right ) :
75+ let leftDepth = left . depth
76+ let rightDepth = right . depth
77+ return 1 + leftDepth > rightDepth ? leftDepth : rightDepth
78+ }
79+
8180 }
82-
83- public var count : Int {
84- get {
85- return Treap . countHelper ( self )
86- }
81+ }
82+
83+ public var count : Int {
84+ get {
85+ return Treap . countHelper ( self )
8786 }
88-
89- private static func countHelper( treap: Treap < Key , Element > ) -> Int {
90- if case let . Node( _, _, _, left, right) = treap {
91- return countHelper ( left) + 1 + countHelper( right)
92- }
93-
94- return 0
87+ }
88+
89+ fileprivate static func countHelper( _ treap: Treap < Key , Element > ) -> Int {
90+ if case let . node( _, _, _, left, right) = treap {
91+ return countHelper ( left) + 1 + countHelper( right)
9592 }
93+
94+ return 0
95+ }
9696}
9797
98- internal func leftRotate< Key: Comparable , Element> ( tree: Treap < Key , Element > ) -> Treap < Key , Element > {
99- if case let . Node ( key, val, p, . Node ( leftKey, leftVal, leftP, leftLeft, leftRight) , right) = tree {
100- return . Node ( key: leftKey, val: leftVal, p: leftP, left: leftLeft,
101- right: Treap . Node ( key: key, val: val, p: p, left: leftRight, right: right) )
102- } else {
103- return . Empty
104- }
98+ internal func leftRotate< Key: Comparable , Element> ( _ tree: Treap < Key , Element > ) -> Treap < Key , Element > {
99+ if case let . node ( key, val, p, . node ( leftKey, leftVal, leftP, leftLeft, leftRight) , right) = tree {
100+ return . node ( key: leftKey, val: leftVal, p: leftP, left: leftLeft,
101+ right: Treap . node ( key: key, val: val, p: p, left: leftRight, right: right) )
102+ } else {
103+ return . empty
104+ }
105105}
106106
107- internal func rightRotate< Key: Comparable , Element> ( tree: Treap < Key , Element > ) -> Treap < Key , Element > {
108- if case let . Node ( key, val, p, left, . Node ( rightKey, rightVal, rightP, rightLeft, rightRight) ) = tree {
109- return . Node ( key: rightKey, val: rightVal, p: rightP,
110- left: Treap . Node ( key: key, val: val, p: p, left: left, right: rightLeft) , right: rightRight)
111- } else {
112- return . Empty
113- }
107+ internal func rightRotate< Key: Comparable , Element> ( _ tree: Treap < Key , Element > ) -> Treap < Key , Element > {
108+ if case let . node ( key, val, p, left, . node ( rightKey, rightVal, rightP, rightLeft, rightRight) ) = tree {
109+ return . node ( key: rightKey, val: rightVal, p: rightP,
110+ left: Treap . node ( key: key, val: val, p: p, left: left, right: rightLeft) , right: rightRight)
111+ } else {
112+ return . empty
113+ }
114114}
115115
116116public extension Treap {
117- internal func set( key: Key , val: Element , p: Int = Int ( arc4random ( ) ) ) -> Treap {
118- switch self {
119- case . Empty:
120- return . Node( key: key, val: val, p: p, left: . Empty, right: . Empty)
121- case let . Node( nodeKey, nodeVal, nodeP, left, right) where key != nodeKey:
122- return insertAndBalance ( nodeKey, nodeVal, nodeP, left, right, key, val, p)
123- case let . Node( nodeKey, _, nodeP, left, right) where key == nodeKey:
124- return . Node( key: key, val: val, p: nodeP, left: left, right: right)
125- default : // should never happen
126- return . Empty
127- }
128-
117+ internal func set( key: Key , val: Element , p: Int = Int ( arc4random ( ) ) ) -> Treap {
118+ switch self {
119+ case . empty:
120+ return . node( key: key, val: val, p: p, left: . empty, right: . empty)
121+ case let . node( nodeKey, nodeVal, nodeP, left, right) where key != nodeKey:
122+ return insertAndBalance ( nodeKey, nodeVal, nodeP, left, right, key, val, p)
123+ case let . node( nodeKey, _, nodeP, left, right) where key == nodeKey:
124+ return . node( key: key, val: val, p: nodeP, left: left, right: right)
125+ default : // should never happen
126+ return . empty
129127 }
130-
131- private func insertAndBalance( nodeKey: Key , _ nodeVal: Element , _ nodeP: Int , _ left: Treap ,
132- _ right: Treap , _ key: Key , _ val: Element , _ p: Int ) -> Treap {
133- let newChild : Treap < Key , Element >
134- let newNode : Treap < Key , Element >
135- let rotate : ( Treap ) -> Treap
136- if key < nodeKey {
137- newChild = left. set ( key, val: val, p: p)
138- newNode = . Node( key: nodeKey, val: nodeVal, p: nodeP, left: newChild, right: right)
139- rotate = leftRotate
140- } else if key > nodeKey {
141- newChild = right. set ( key, val: val, p: p)
142- newNode = . Node( key: nodeKey, val: nodeVal, p: nodeP, left: left, right: newChild)
143- rotate = rightRotate
144- } else {
145- // It should be impossible to reach here
146- newChild = . Empty
147- newNode = . Empty
148- return newNode
149- }
150-
151- if case let . Node( _, _, newChildP, _, _) = newChild where newChildP < nodeP {
152- return rotate ( newNode)
153- } else {
154- return newNode
155- }
128+ }
129+
130+ fileprivate func insertAndBalance( _ nodeKey: Key , _ nodeVal: Element , _ nodeP: Int , _ left: Treap ,
131+ _ right: Treap , _ key: Key , _ val: Element , _ p: Int ) -> Treap {
132+ let newChild : Treap < Key , Element >
133+ let newNode : Treap < Key , Element >
134+ let rotate : ( Treap ) -> Treap
135+ if key < nodeKey {
136+ newChild = left. set ( key: key, val: val, p: p)
137+ newNode = . node( key: nodeKey, val: nodeVal, p: nodeP, left: newChild, right: right)
138+ rotate = leftRotate
139+ } else if key > nodeKey {
140+ newChild = right. set ( key: key, val: val, p: p)
141+ newNode = . node( key: nodeKey, val: nodeVal, p: nodeP, left: left, right: newChild)
142+ rotate = rightRotate
143+ } else {
144+ // It should be impossible to reach here
145+ newChild = . empty
146+ newNode = . empty
147+ return newNode
156148 }
157-
158- internal func delete( key: Key ) throws -> Treap {
159- switch self {
160- case . Empty:
161- throw NSError ( domain: " com.wta.treap.errorDomain " , code: - 1 , userInfo: nil )
162- case let . Node( nodeKey, val, p, left, right) where key < nodeKey:
163- return try Treap . Node ( key: nodeKey, val: val, p: p, left: left. delete ( key) , right: right)
164- case let . Node( nodeKey, val, p, left, right) where key > nodeKey:
165- return try Treap . Node ( key: nodeKey, val: val, p: p, left: left, right: right. delete ( key) )
166- case let . Node( _, _, _, left, right) :
167- return merge ( left, right: right)
168- }
149+
150+ if case let . node( _, _, newChildP, _, _) = newChild , newChildP < nodeP {
151+ return rotate ( newNode)
152+ } else {
153+ return newNode
154+ }
155+ }
156+
157+ internal func delete( key: Key ) throws -> Treap {
158+ switch self {
159+ case . empty:
160+ throw NSError ( domain: " com.wta.treap.errorDomain " , code: - 1 , userInfo: nil )
161+ case let . node( nodeKey, val, p, left, right) where key < nodeKey:
162+ return try Treap . node ( key: nodeKey, val: val, p: p, left: left. delete ( key: key) , right: right)
163+ case let . node( nodeKey, val, p, left, right) where key > nodeKey:
164+ return try Treap . node ( key: nodeKey, val: val, p: p, left: left, right: right. delete ( key: key) )
165+ case let . node( _, _, _, left, right) :
166+ return merge ( left, right: right)
169167 }
168+ }
170169}
0 commit comments