Untitled

mail@pastecode.io avatar
unknown
plain_text
a year ago
1.6 kB
6
Indexable
Never
from queue import Queue
from heapq import nlargest
from typing import List

class City:
    def __init__(self):
        self.restaurants = []
        self.adjList = []

city = []
value = {}
best = {}

def init(N, M, mRoads):
    global city, value, best
    value = {}
    best = {}
    city = [City() for _ in range(N + 1)]

    for i in range(M):
        city[mRoads[i][0]].adjList.append(mRoads[i][1])
        city[mRoads[i][1]].adjList.append(mRoads[i][0])

def addRestaurant(mCityID, mName):
    city[mCityID].restaurants.append(''.join(mName))

def addValue(mName, mScore):
    global value, best
    name = ''.join(mName)
    score = value.get(name, 0) + mScore

    value[name] = score

    for i in range(len(name)):
        for j in range(1, len(name) - i + 1):
            sub = name[i:i + j]
            if best.get(sub, 0) < score:
                best[sub] = score

def bestValue(mStr):
    return best.get(''.join(mStr), 0)

def regionalValue(mCityID, mDist):
    global city, value
    queue = Queue()
    top3 = []
    dist = [0] * len(city)
    queue.put(mCityID)

    while not queue.empty():
        currCity = queue.get()

        for name in city[currCity].restaurants:
            top3.append(value.get(name, 0))
            if len(top3) > 3:
                top3.remove(min(top3))

        for i in city[currCity].adjList:
            if dist[i] == 0 and i != mCityID:
                dist[i] = dist[currCity] + 1
                if dist[i] <= mDist:
                    queue.put(i)

    return sum(top3)