From 34afb450d81b1002a897d7658736747bf929b910 Mon Sep 17 00:00:00 2001 From: Ilya Gorbunov Date: Tue, 21 Jul 2015 20:03:56 +0300 Subject: [PATCH] List binarySearch and binarySearchBy methods and tests. #KT-5444 Fixed #KT-8217 Fixed --- .../stdlib/src/kotlin/collections/JUtil.kt | 106 ++++++++++++++++++ .../test/collections/ListBinarySearchTest.kt | 106 ++++++++++++++++++ 2 files changed, 212 insertions(+) create mode 100644 libraries/stdlib/test/collections/ListBinarySearchTest.kt diff --git a/libraries/stdlib/src/kotlin/collections/JUtil.kt b/libraries/stdlib/src/kotlin/collections/JUtil.kt index 572ac9ee6f8..65eeb126379 100644 --- a/libraries/stdlib/src/kotlin/collections/JUtil.kt +++ b/libraries/stdlib/src/kotlin/collections/JUtil.kt @@ -156,3 +156,109 @@ private fun Iterable.convertToSetForSetOperation(): Collection = is Collection -> if (this.safeToConvertToSet()) toHashSet() else this else -> toHashSet() } + +/** + * Searches this list or its range for the provided [element] index using binary search algorithm. + * The list is expected to be sorted into ascending order according to the Comparable natural ordering of its elements. + * + * If the list contains multiple elements equal to the specified object, there is no guarantee which one will be found. + */ +public fun > List.binarySearch(element: T?, fromIndex: Int = 0, toIndex: Int = size()): Int { + rangeCheck(size(), fromIndex, toIndex) + + var low = fromIndex + var high = toIndex - 1 + + while (low <= high) { + val mid = (low + high).ushr(1) // safe from overflows + val midVal = get(mid) + val cmp = compareValues(midVal, element) + + if (cmp < 0) + low = mid + 1 + else if (cmp > 0) + high = mid - 1 + else + return mid // key found + } + return -(low + 1) // key not found +} + +/** + * Searches this list or its range for the provided [element] index using binary search algorithm. + * The list is expected to be sorted into ascending order according to the specified [comparator]. + * + * If the list contains multiple elements equal to the specified object, there is no guarantee which one will be found. + */ +public fun List.binarySearch(element: T, comparator: Comparator, fromIndex: Int = 0, toIndex: Int = size()): Int { + rangeCheck(size(), fromIndex, toIndex) + + var low = fromIndex + var high = toIndex - 1 + + while (low <= high) { + val mid = (low + high).ushr(1) // safe from overflows + val midVal = get(mid) + val cmp = comparator.compare(midVal, element) + + if (cmp < 0) + low = mid + 1 + else if (cmp > 0) + high = mid - 1 + else + return mid // key found + } + return -(low + 1) // key not found +} + +/** + * Searches this list or its range for an index of an element with the provided [key] using binary search algorithm. + * The list is expected to be sorted into ascending order according to the Comparable natural ordering of keys of its elements. + * + * If the list contains multiple elements with the specified [key], there is no guarantee which one will be found. + */ +public inline fun > List.binarySearchBy(key: K?, fromIndex: Int = 0, toIndex: Int = size(), inlineOptions(InlineOption.ONLY_LOCAL_RETURN) selector: (T) -> K?): Int = + binarySearch(fromIndex, toIndex) { compareValues(selector(it), key) } + +// do not introduce this overload --- too rare +//public fun List.binarySearchBy(key: K, comparator: Comparator, fromIndex: Int = 0, toIndex: Int = size(), selector: (T) -> K): Int = +// binarySearch(fromIndex, toIndex) { comparator.compare(selector(it), key) } + +/** + * Searches this list or its range for an index of an element for which [comparison] function returns zero. + * The list is expected to be sorted into ascending order according to the provided [comparison]. + * + * @param comparison function that compares an element of the list with the element being searched. + */ +public fun List.binarySearch(fromIndex: Int = 0, toIndex: Int = size(), comparison: (T) -> Int): Int { + rangeCheck(size(), fromIndex, toIndex) + + var low = fromIndex + var high = toIndex - 1 + + while (low <= high) { + val mid = (low + high).ushr(1) // safe from overflows + val midVal = get(mid) + val cmp = comparison(midVal) + + if (cmp < 0) + low = mid + 1 + else if (cmp > 0) + high = mid - 1 + else + return mid // key found + } + return -(low + 1) // key not found +} + +/** + * Checks that `from` and `to` are in + * the range of [0..size] and throws an appropriate exception, if they aren't. + */ +private fun rangeCheck(size: Int, fromIndex: Int, toIndex: Int) { + when { + fromIndex > toIndex -> throw IllegalArgumentException("fromIndex ($fromIndex) is greater than toIndex ($toIndex)") + fromIndex < 0 -> throw IndexOutOfBoundsException("fromIndex ($fromIndex) is less than zero.") + toIndex > size -> throw IndexOutOfBoundsException("toIndex ($toIndex) is greater than size ($size).") + } +} diff --git a/libraries/stdlib/test/collections/ListBinarySearchTest.kt b/libraries/stdlib/test/collections/ListBinarySearchTest.kt new file mode 100644 index 00000000000..6deff32fd74 --- /dev/null +++ b/libraries/stdlib/test/collections/ListBinarySearchTest.kt @@ -0,0 +1,106 @@ +package test.collections.binarySearch + +import org.junit.Test +import kotlin.test.assertEquals + +class ListBinarySearchTest { + + val values = listOf(1, 3, 7, 10, 12, 15, 22, 45) + + fun notFound(index: Int) = -(index + 1) + + val comparator = compareBy?> { it?.value } + + Test fun binarySearchByElement() { + val list = values + list.forEachIndexed { index, item -> + assertEquals(index, list.binarySearch(item)) + assertEquals(notFound(index), list.binarySearch(item.pred())) + assertEquals(notFound(index + 1), list.binarySearch(item.succ())) + + if (index > 0) { + index.let { from -> assertEquals(notFound(from), list.binarySearch(list.first(), fromIndex = from)) } + (list.size() - index).let { to -> assertEquals(notFound(to), list.binarySearch(list.last(), toIndex = to)) } + } + } + } + + Test fun binarySearchByElementNullable() { + val list = listOf(null) + values + list.forEachIndexed { index, item -> + assertEquals(index, list.binarySearch(item)) + + if (index > 0) { + index.let { from -> assertEquals(notFound(from), list.binarySearch(list.first(), fromIndex = from)) } + (list.size() - index).let { to -> assertEquals(notFound(to), list.binarySearch(list.last(), toIndex = to)) } + } + } + } + + Test fun binarySearchWithComparator() { + val list = values map { IncomparableDataItem(it) } + + list.forEachIndexed { index, item -> + assertEquals(index, list.binarySearch(item, comparator)) + assertEquals(notFound(index), list.binarySearch(item.pred(), comparator)) + assertEquals(notFound(index + 1), list.binarySearch(item.succ(), comparator)) + + if (index > 0) { + index.let { from -> assertEquals(notFound(from), list.binarySearch(list.first(), comparator, fromIndex = from)) } + (list.size() - index).let { to -> assertEquals(notFound(to), list.binarySearch(list.last(), comparator, toIndex = to)) } + } + } + } + + Test fun binarySearchByKey() { + val list = values map { IncomparableDataItem(it) } + + list.forEachIndexed { index, item -> + assertEquals(index, list.binarySearchBy(item.value) { it.value }) + assertEquals(notFound(index), list.binarySearchBy(item.value.pred()) { it.value }) + assertEquals(notFound(index + 1), list.binarySearchBy(item.value.succ()) { it.value }) + + if (index > 0) { + index.let { from -> assertEquals(notFound(from), list.binarySearchBy(list.first().value, fromIndex = from) { it.value }) } + (list.size() - index).let { to -> assertEquals(notFound(to), list.binarySearchBy(list.last().value, toIndex = to) { it.value }) } + } + } + } + + + Test fun binarySearchByKeyWithComparator() { + val list = values map { IncomparableDataItem(IncomparableDataItem(it)) } + + list.forEachIndexed { index, item -> + assertEquals(index, list.binarySearch { comparator.compare(it.value, item.value) } ) + assertEquals(notFound(index), list.binarySearch { comparator.compare(it.value, item.value.pred()) }) + assertEquals(notFound(index + 1), list.binarySearch { comparator.compare(it.value, item.value.succ()) }) + + if (index > 0) { + index.let { from -> + assertEquals(notFound(from), list.binarySearch(fromIndex = from) { comparator.compare(it.value, list.first().value) }) + } + (list.size() - index).let { to -> + assertEquals(notFound(to), list.binarySearch(toIndex = to) { comparator.compare(it.value, list.last().value) }) + } + } + } + } + + Test fun binarySearchByMultipleKeys() { + val list = values.flatMap { v1 -> values.map { v2 -> Pair(v1, v2) } } + + list.forEachIndexed { index, item -> + assertEquals(index, list.binarySearch { compareValuesBy(it, item, { it.first }, { it.second }) }) + } + } +} + + +private data class IncomparableDataItem(public val value: T) +private fun IncomparableDataItem.pred(): IncomparableDataItem = IncomparableDataItem(value - 1) +private fun IncomparableDataItem.succ(): IncomparableDataItem = IncomparableDataItem(value + 1) +private fun Int.pred() = dec() +private fun Int.succ() = inc() + +