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")