https://coderun.yandex.ru/problem/diversity-improvement Сложная

Решение

import java.io.BufferedReader
import java.io.BufferedWriter
import java.io.InputStreamReader
import java.io.OutputStreamWriter
import java.util.PriorityQueue
import java.util.ArrayDeque

data class Product(val id: Int, val category: Int)

fun main(args: Array<String>) {
    val reader = BufferedReader(InputStreamReader(System.`in`))
    val writer = BufferedWriter(OutputStreamWriter(System.out))
    
    // Чтение входных данных
    val n = reader.readLine().toInt()
    val productsInput = ArrayList<Product>(n)
    // Группируем товары по категориям и считаем количество
    val catCounts = HashMap<Int, Int>()
    val catProducts = HashMap<Int, ArrayDeque<Int>>()
    for (i in 0 until n) {
        val parts = reader.readLine().split(" ")
        val prodId = parts[0].toInt()
        val cat = parts[1].toInt()
        productsInput.add(Product(prodId, cat))
        catCounts[cat] = (catCounts[cat] ?: 0) + 1
        if (!catProducts.containsKey(cat)) {
            catProducts[cat] = ArrayDeque()
        }
        catProducts[cat]!!.add(prodId)
    }
    
    // Если все товары уникальны по категориям – выводим исходный порядок
    if (catCounts.values.all { it == 1 }) {
        for (p in productsInput) {
            writer.write("${p.id} ")
        }
        writer.flush()
        reader.close()
        writer.close()
        return
    }
    
    // Находим максимальную частоту и число категорий, для которых эта частота достигается
    var maxFreq = 0
    for (cnt in catCounts.values) {
        if (cnt > maxFreq) {
            maxFreq = cnt
        }
    }
    var countMax = 0
    for (cnt in catCounts.values) {
        if (cnt == maxFreq) {
            countMax++
        }
    }
    
    // Вычисляем оптимальное минимальное расстояние d
    val d = (n - countMax) / (maxFreq - 1)
    
    // Опишем задачу для симуляции: для каждой категории будем отслеживать
    // оставшееся число товаров и время, когда можно вставить следующий товар
    data class Task(val cat: Int, var remaining: Int, var nextAvailable: Int)
    
    // Очередь доступных задач: сортируем по оставшемуся количеству (по убыванию)
    val available = PriorityQueue<Task> { a, b -> b.remaining - a.remaining }
    // Очередь ожидающих задач: сортировка по времени, когда задача станет доступной (по возрастанию)
    val waiting = PriorityQueue<Task> { a, b -> a.nextAvailable - b.nextAvailable }
    
    // Инициализируем все задачи. В начале все доступны (nextAvailable = 0).
    for ((cat, cnt) in catCounts) {
        available.add(Task(cat, cnt, 0))
    }
    
    val result = IntArray(n)
    // Симуляция заполнения позиций от 0 до n-1
    for (i in 0 until n) {
        // «Размораживаем» задачи, время доступности которых истекло
        while (waiting.isNotEmpty() && waiting.peek().nextAvailable <= i) {
            available.add(waiting.poll())
        }
        if (available.isEmpty()) {
            // По теории это происходить не должно, поскольку d вычислено оптимально
            writer.write("Нет решения")
            writer.flush()
            reader.close()
            writer.close()
            return
        }
        // Выбираем задачу с максимальным оставшимся числом товаров
        val task = available.poll()
        // Из соответствующей очереди товаров извлекаем один product id
        val prodId = catProducts[task.cat]!!.removeFirst()
        result[i] = prodId
        task.remaining--
        // Если в категории ещё есть товары, обновляем время доступности и помещаем в очередь ожидания.
        if (task.remaining > 0) {
            task.nextAvailable = i + d
            waiting.add(task)
        }
    }
    
    // Вывод результата – перестановку product id
    for (id in result) {
        writer.write("$id ")
    }
    
    writer.flush()
    reader.close()
    writer.close()
}