Untitled
unknown
python
a year ago
2.0 kB
4
Indexable
Never
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")