search system

import java.lang.Integer.min
import java.util.*
import kotlin.collections.HashMap

class SearchSystem(
    documents: Array<String>
) {
    // в value хранится номер документа и количество вхождений слова в этом документе
    // word to { docNumber to wordCountInDocument }
    private val wordStats: MutableMap<String, MutableMap<Int, Int>> = HashMap(documents.size)

    init {
        for (docNumber in 1..documents.size) {
            documents[docNumber - 1].asTokenSequence().forEach {
                    if (wordStats[it] == null) {
                        wordStats[it] = mutableMapOf(docNumber to 1)
                    } else {
                        wordStats[it]!!.merge(docNumber, 1) { wordCount, _ -> wordCount + 1 }
                    }
                }
        }
    }

    fun findDocumentNumbers(requests: Array<String>, maxCount: Int): List<List<Int>> {
        return requests.asSequence()
            .map { request -> searchByUniqueWords(request.asTokenSequence().toSet()).entries.toMutableList()
                .takeTop(maxCount)
                .map { it.key }
            }
            .filter { it.isNotEmpty() }
            .toList()
    }

    private fun searchByUniqueWords(uniqueWords: Set<String>): Map<Int, Int> {
        return wordStats.asSequence()
            .filter { uniqueWords.contains(it.key) }                    // данные только по словам из запроса
            .flatMap { it.value.entries }                               // статистика только по словам из запроса
            .groupBy { it.key }                                         // группировка по номеру документа
            .mapValues { entry -> entry.value.sumOf { it.value } }      // подсчёт суммы вхождений в документе
    }

    private operator fun Map.Entry<Int,Int>.compareTo(other: Map.Entry<Int,Int>): Int {
        val countDifference = other.value - this.value
        return if (countDifference == 0) {
            this.key - other.key
        } else {
            countDifference
        }
    }

    private fun MutableList<Map.Entry<Int,Int>>.takeTop(maxCount: Int): List<Map.Entry<Int,Int>>{
        if (this.size < 2) {
            return this
        }
        val maxSize = min(this.size, maxCount)
        for (i in 0 until maxSize) {
            var maxItem = this[i]
            var j = i
            while (j < this.size) {
                if (this[j] < maxItem) {
                    val temp = this[j]
                    this[j] = maxItem
                    maxItem = temp
                }
                j++
            }
            this[i] = maxItem
        }

        return this.subList(0, maxSize)
    }

    private fun String.asTokenSequence(): Sequence<String> {
        val stringTokenizer = StringTokenizer(this)
        return generateSequence {
            if (stringTokenizer.hasMoreTokens()) {
                stringTokenizer.nextToken()
            } else {
                null
            }
        }
    }
}

fun main() {
    val reader = System.`in`.bufferedReader()
    val writer = System.out.bufferedWriter()
    val documentsCount = reader.readLine().toInt()
    val documents = Array(documentsCount) {
        reader.readLine()
    }
    val requestsCount = reader.readLine().toInt()
    val requests = Array(requestsCount) {
        reader.readLine()
    }

    val searchSystem = SearchSystem(documents)
    val results = searchSystem.findDocumentNumbers(requests, 5)
    with(writer) {
        results.forEach {
                write(it.joinToString(" "))
                newLine()
                flush()
            }
    }
}