solution.cpp

#include <vector>
#include <queue>
#include <unordered_map>
#include <cstdio>

using namespace std;

// -------- structures --------

// -------- global variables --------

int n;
vector<tuple<int, int, int>> graph[300];  // e, cost, mId

unordered_map<int, pair<int, int>> road_location; // mId, s, index

unordered_map<int, int> map_rid;
int map_rid_len;

// -------- helper functions --------

int get_rid(const int mId) {
  auto it = map_rid.find(mId);
  if (it == map_rid.end()) {
    return map_rid[mId] = map_rid_len++;
  }

  return it->second;
}

void add_road(const int mId, const int s, const int e, const int cost) {
  // add road
  const int rid = get_rid(mId);
  graph[s].push_back({e, cost, mId});

  // add road location info
  const int index = graph[s].size() - 1;
  road_location[mId] = {s, index};
}

// -------- API --------

void init(int N, int K, int mId[], int sCity[], int eCity[], int mToll[]) {
  n = N;

  for (int i = 0; i < n; i++) {
    graph[i].clear();
  }

  road_location.clear();
  map_rid.clear();
  map_rid_len = 0;

  for (int i = 0; i < K; i++) {
    add_road(mId[i], sCity[i], eCity[i], mToll[i]);
  }
}

// 1,400
void add(int mId, int sCity, int eCity, int mToll) {
  add_road(mId, sCity, eCity, mToll);
}

// 500
void remove(int mId) {
  const auto location = road_location[mId];

  const int s = location.first;
  const int index = location.second;

  graph[s][index] = *graph[s].rbegin();
  graph[s].pop_back();

  if (index < graph[s].size()) {
    const int mId2 = get<2>(graph[s][index]);
    road_location[mId2].second = index;
  }
}

// 100
int cost(int M, int sCity, int eCity) {
  priority_queue<tuple<int, int, int>, vector<tuple<int, int, int>>, greater<tuple<int, int, int>>> pq; // cost, m, node

  pq.push({0, M, sCity});

  while (pq.size()) {
    int cc, cm, cn;
    tie(cc, cm, cn) = pq.top();  // cost, m, node
    pq.pop();

    if (cn == eCity) {
      return cc;
    }

    for (int i = 0; i < graph[cn].size(); i++) {
      int nn, cost, _;
      tie(nn, cost, _) = graph[cn][i];

      pq.push({cc + cost, cm, nn});
      if (cm > 0) {
        pq.push({cc + cost / 2, cm - 1, nn});
      }
    }
  }

  return -1;
}