FlowerDelivery.kt

package com.revolut.interview

import java.util.Collections.emptyList

data class Flower(
    val id: Int,
    val color: String
)

data class Address(
    val address: String
)

data class Box(
    val address: Address,
    val flowers: List<Flower> = emptyList()
)

class FlowerDelivery {
    /*
     * Distributes flowers over given addresses aligning the following rules:
     * 1) Each address should get a box
     * 2) Each box should contain at least one flower
     * 3) If box contains more than one flower then flower colours should be unique
     */
    // Time complexity: O(N + M)
    fun distribute(addresses: List<Address>, flowers: List<Flower>): List<Box> { // N == addresses.size, M == flowers.size
        if (addresses.isEmpty()) {
            throw IllegalStateException("Addresses must be defined")
        }

        if (flowers.isEmpty()) {
            throw IllegalStateException("Flowers must present")
        }

        if (flowers.size < addresses.size) {
            throw IllegalArgumentException("Cannot distribute ${flowers.size} flower for ${addresses.size} addresses")
        }

        // O(N * max len of address), but in real life I would distinct by `id` so it would be O(N)
        val uniqueAddresses = addresses.distinctBy { it.address }
        if (uniqueAddresses.size < addresses.size) {
            val notUnique: List<Address> = (addresses - uniqueAddresses) // O(N) but nevermind because it's an error branch - doing it for debug purposes
            throw IllegalStateException("Given addresses presented in list twice: $notUnique")
        }

        val addressFlowers: List<MutableList<Flower>> = ArrayList<MutableList<Flower>>(addresses.size).apply {
            repeat(addresses.size) { // O(N)
                add(mutableListOf())
            }
        }

        val flowersByColor: Map<String, List<Flower>> = flowers.groupBy { it.color } // O(M)

        var currentAddressIndex = 0
        for ((_, colorFlowers) in flowersByColor) { // O(const) == O(1) because there's no much colors
            val startIndex = currentAddressIndex
            for (colorFlower in colorFlowers) { // O(M/const) == O(M)
                addressFlowers[currentAddressIndex].add(colorFlower)

                currentAddressIndex = (currentAddressIndex + 1) % addresses.size
                if (currentAddressIndex == startIndex) { // We distributed as much flowers for this color as we could
                    break
                }
            }
        }

        val boxes: List<Box> = addressFlowers.mapIndexed { index, boxFlowers -> // O(N)
            Box(addresses[index], boxFlowers)
        }

        // Final complexity:
        // O(N + N + M + M + N) = O(3N + 2M) = O(N + M)

        return boxes
    }
}