diff --git a/Sources/RedBlackTree.swift b/Sources/RedBlackTree.swift index eaf259a..6db6e03 100644 --- a/Sources/RedBlackTree.swift +++ b/Sources/RedBlackTree.swift @@ -348,6 +348,24 @@ public struct RedBlackTree: Probable, Collection, Bidire return internalFindNodeForKey(key).value } + /** + :name: findLowerValue + :description: Finds instance with key that is lower or equal input key + - returns: Value? + */ + public func findLowerValue(for key: Key) -> Value? { + return internalFindLowerForKey(key).value + } + + /** + :name: findCeilingValue + :description: Finds instance with key that is larger or equal input key + - returns: Value? + */ + public func findCeilingValue(for key: Key) -> Value? { + return internalFindCeilingForKey(key).value + } + /** Returns the Key value at a given position. - Parameter position: An Int. @@ -395,6 +413,48 @@ public struct RedBlackTree: Probable, Collection, Bidire } } + /** + :name: internalFindLowerForKey + :description: Finds a node with a key that is equal or less that given. + - returns: RedBlackNode + */ + private func internalFindLowerForKey(_ key: Key) -> RedBlackNode { + var z = root + var max = sentinel + + while z !== sentinel { + if key > z.key { + max = z + } + if key == z.key { + return z + } + z = key < z.key as Key ? z.left : z.right + } + return max + } + + /** + :name: internalFindCeilingForKey + :description: Finds a node with a key that is equal or larger that given. + - returns: RedBlackNode + */ + private func internalFindCeilingForKey(_ key: Key) -> RedBlackNode { + var z = root + var min = sentinel + + while z !== sentinel { + if key < z.key { + min = z + } + if key == z.key { + return z + } + z = key < z.key as Key ? z.left : z.right + } + return min + } + /** :name: indexOf :description: Returns the Index of a given member, or nil if the member is not present in the set. diff --git a/Sources/SortedDictionary.swift b/Sources/SortedDictionary.swift index f3a60e5..23824f0 100644 --- a/Sources/SortedDictionary.swift +++ b/Sources/SortedDictionary.swift @@ -309,6 +309,24 @@ public struct SortedDictionary: Probable, Col return tree.findValue(for: key) } + /** + Finds the value for a key which is less or equal given. + - Parameter for key: A Key type. + - Returns: An optional Value type. + */ + public func findLowerEntry(for key: Key) -> Value? { + return tree.findLowerValue(for: key) + } + + /** + Finds the value for a key which is larger or equal given. + - Parameter for key: A Key type. + - Returns: An optional Value type. + */ + public func findCeilingEntry(for key: Key) -> Value? { + return tree.findCeilingValue(for: key) + } + /** Searches for given keys in the SortedDictionary. - Parameter for keys: A list of Key types. diff --git a/Tests/SortedDictionaryTests.swift b/Tests/SortedDictionaryTests.swift index 8e9ac13..2517cc7 100644 --- a/Tests/SortedDictionaryTests.swift +++ b/Tests/SortedDictionaryTests.swift @@ -121,4 +121,42 @@ class SortedDictionaryTests: XCTestCase { let values = [1, 2, 3, 4] XCTAssert(values == s.values, "Test failed.") } + + func testLowerEntry() { + let s = SortedDictionary(elements: (1, 1), (2, 2), (3, 3), (5, 5), (8, 8), (13, 13), (21, 21), (34, 34)) + + XCTAssert(s.findLowerEntry(for: -15) == nil, "Test failed.") + XCTAssert(s.findLowerEntry(for: 0) == nil, "Test failed.") + XCTAssert(s.findLowerEntry(for: 1) == 1, "Test failed.") + XCTAssert(s.findLowerEntry(for: 2) == 2, "Test failed.") + XCTAssert(s.findLowerEntry(for: 3) == 3, "Test failed.") + XCTAssert(s.findLowerEntry(for: 4) == 3, "Test failed.") + XCTAssert(s.findLowerEntry(for: 5) == 5, "Test failed.") + XCTAssert(s.findLowerEntry(for: 6) == 5, "Test failed.") + XCTAssert(s.findLowerEntry(for: 7) == 5, "Test failed.") + XCTAssert(s.findLowerEntry(for: 8) == 8, "Test failed.") + XCTAssert(s.findLowerEntry(for: 9) == 8, "Test failed.") + XCTAssert(s.findLowerEntry(for: 10) == 8, "Test failed.") + XCTAssert(s.findLowerEntry(for: 40) == 34, "Test failed.") + XCTAssert(s.findLowerEntry(for: 50) == 34, "Test failed.") + } + + func testCeilingEntry() { + let s = SortedDictionary(elements: (1, 1), (2, 2), (3, 3), (5, 5), (8, 8), (13, 13), (21, 21), (34, 34)) + + XCTAssert(s.findCeilingEntry(for: -15) == 1, "Test failed.") + XCTAssert(s.findCeilingEntry(for: 0) == 1, "Test failed.") + XCTAssert(s.findCeilingEntry(for: 1) == 1, "Test failed.") + XCTAssert(s.findCeilingEntry(for: 2) == 2, "Test failed.") + XCTAssert(s.findCeilingEntry(for: 3) == 3, "Test failed.") + XCTAssert(s.findCeilingEntry(for: 4) == 5, "Test failed.") + XCTAssert(s.findCeilingEntry(for: 5) == 5, "Test failed.") + XCTAssert(s.findCeilingEntry(for: 6) == 8, "Test failed.") + XCTAssert(s.findCeilingEntry(for: 7) == 8, "Test failed.") + XCTAssert(s.findCeilingEntry(for: 8) == 8, "Test failed.") + XCTAssert(s.findCeilingEntry(for: 9) == 13, "Test failed.") + XCTAssert(s.findCeilingEntry(for: 10) == 13, "Test failed.") + XCTAssert(s.findCeilingEntry(for: 40) == nil, "Test failed.") + XCTAssert(s.findCeilingEntry(for: 50) == nil, "Test failed.") + } }