Kotlin如何不推进迭代器查看下一个元素？无需Guava优化合并序列

当然可以在不推进迭代器的前提下查看下一个元素，Kotlin标准库没有内置PeekingIterator，但我们可以通过自定义封装或利用序列(Sequence)的特性来实现，完全不需要依赖第三方库。

先看原始的繁琐实现（无Peek功能） #

如果直接用标准Iterator合并两个预排序序列，你不得不每次先判断元素、取出后还要手动维护临时变量（因为next()会直接推进迭代器），代码会变得冗余且不直观：

fun <T : Comparable<T>> mergeSorted(a: Iterator<T>, b: Iterator<T>): Sequence<T> = sequence {
    var currentA: T? = if (a.hasNext()) a.next() else null
    var currentB: T? = if (b.hasNext()) b.next() else null

    while (currentA != null && currentB != null) {
        if (currentA <= currentB) {
            yield(currentA)
            currentA = if (a.hasNext()) a.next() else null
        } else {
            yield(currentB)
            currentB = if (b.hasNext()) b.next() else null
        }
    }

    // 处理剩余元素
    currentA?.let { yield(it); yieldAll(a) }
    currentB?.let { yield(it); yieldAll(b) }
}

方案1：自定义PeekingIterator类 #

我们可以封装一个自己的PeekingIterator，实现peek()方法来查看下一个元素而不推进迭代器：

class PeekingIterator<T>(private val iterator: Iterator<T>) : Iterator<T> {
    private var cachedNext: T? = null
    private var hasCached = false

    override fun hasNext(): Boolean = hasCached || iterator.hasNext()

    override fun next(): T {
        if (!hasCached) {
            cachedNext = iterator.next()
            hasCached = true
        }
        val result = cachedNext!!
        cachedNext = null
        hasCached = false
        return result
    }

    // 查看下一个元素，不推进迭代器
    fun peek(): T? {
        if (!hasCached) {
            cachedNext = if (iterator.hasNext()) iterator.next() else null
            hasCached = true
        }
        return cachedNext
    }
}

用这个类来合并序列，逻辑会简洁很多，完全贴合合并有序序列的直觉：

fun <T : Comparable<T>> mergeSorted(a: Iterator<T>, b: Iterator<T>): Sequence<T> = sequence {
    val peekA = PeekingIterator(a)
    val peekB = PeekingIterator(b)

    while (peekA.hasNext() && peekB.hasNext()) {
        val nextA = peekA.peek()!!
        val nextB = peekB.peek()!!
        if (nextA <= nextB) {
            yield(peekA.next())
        } else {
            yield(peekB.next())
        }
    }

    // 直接处理剩余元素
    yieldAll(peekA)
    yieldAll(peekB)
}

方案2：利用Kotlin Sequence的状态维护 #

如果不想自定义类，也可以通过扩展函数结合状态变量来实现类似peek的效果：
首先定义一个迭代器的扩展函数，用来安全获取下一个元素（为空则返回null）：

fun <T> Iterator<T>.nextOrNull(): T? = if (hasNext()) next() else null

然后基于这个扩展函数实现合并逻辑：

fun <T : Comparable<T>> mergeSorted(a: Iterator<T>, b: Iterator<T>): Sequence<T> = sequence {
    val iterA = a.iterator()
    val iterB = b.iterator()
    var nextA = iterA.nextOrNull()
    var nextB = iterB.nextOrNull()

    while (nextA != null && nextB != null) {
        if (nextA <= nextB) {
            yield(nextA)
            nextA = iterA.nextOrNull()
        } else {
            yield(nextB)
            nextB = iterB.nextOrNull()
        }
    }

    // 输出剩余元素
    nextA?.let { yield(it); yieldAll(iterA) }
    nextB?.let { yield(it); yieldAll(iterB) }
}

这种方式不需要额外封装类，直接利用Kotlin的扩展函数和序列特性，代码同样简洁直观。

内容的提问来源于stack exchange，提问作者k314159