Untitled

from dataclasses import dataclass
from typing import Set, List
from datetime import date
import random
from statistics import mean
from heapq import heappush, heappop
import time


@dataclass(frozen=True)
class City:
    id: str
    name: str
    population: int


@dataclass
class DailyTemp:
    date: date
    temperature: float


class WeatherAPI:
    def getAllCitiesByIds(self, city_ids: Set[str]) -> Set[City]:
        return {City(city_id, f"City{city_id}", random.randint(10000, 10000000)) for city_id in city_ids}

    def getLastYearTemperature(self, city_id: str) -> List[DailyTemp]:
        return [DailyTemp(date(2023, 1, 1), random.randint(0, 50)) for _ in range(365)]


class TempAggregator:
    def __init__(self, api: WeatherAPI):
        self.api = api

    def get_top_cities(self, city_ids: Set[str], agg_type: str, n: int) -> List[dict]:
        heap = []
        cities = self.api.getAllCitiesByIds(city_ids)
        for city in cities:
            if city.population > 50000:
                daily_temps = self.api.getLastYearTemperature(city.id)
                if daily_temps:
                    temperatures = [temp.temperature for temp in daily_temps]
                    if agg_type == 'avg':
                        aggregated_temp = mean(temperatures)
                    else:
                        raise ValueError("Aggregation type unknown")

                    heappush(heap, (-aggregated_temp, city.name))
        top_cities = []

        while heap and n > 0:
            temp, city = heappop(heap)
            top_cities.append({city: -temp})
            n -= 1

        return top_cities


N = 5
aggregation_type = 'avg'

weather_api = WeatherAPI()
aggregator = TempAggregator(weather_api)
time1 = time.time()
city_ids = {f"C{i}" for i in range(50000)}
top_cities = aggregator.get_top_cities(city_ids, aggregation_type, N)

for i, city in enumerate(top_cities, start=1):
    print(f"#{i}. {city}")
t = time.time() - time1
print(f"{t:4f}s")