Problem p02377 (C++) — buggy submission, verdict: Time Limit Exceeded


#include <algorithm>
#include <iostream>
#include <limits>
#include <map>
#include <queue>
#include <set>
#include <stdio.h>
#include <time.h>
#include <utility>
#include <vector>

using namespace std;

class Edge {
public:
  int from;
  int to;
  int flow;
  Edge() {}
  Edge(int f, int t, int d) {
    from = f;
    to = t;
    flow = d;
  }
  bool operator==(const Edge &e) {
    return this->from == e.from && this->to == e.to;
  }
};

class Shortest_path_result {
public:
  int sum_of_cost;
  vector<int> path;
  vector<int> distance;
  vector<int> predecessor;
  Shortest_path_result() {}
};

class Graph {
public:
  int INF;
  int n;
  vector<set<int>> vertices_list;
  vector<map<int, int>> cost_list;
  vector<map<int, int>> capacity_list;
  vector<int> potential_list;

  Graph() {}
  Graph(int n) { // nは頂点数
    INF = 2e9;
    this->n = n;
    // 各種リストを集合
    // n個の空集合を挿入
    vertices_list.insert(vertices_list.begin(), n, set<int>());
    cost_list.insert(cost_list.begin(), n, map<int, int>());
    capacity_list.insert(capacity_list.begin(), n, map<int, int>());
    potential_list = vector<int>(n, 0);
  }

  Graph clone(void) {
    Graph g = Graph(this->n);
    g.vertices_list = this->vertices_list;
    g.cost_list = this->cost_list;
    g.capacity_list = this->capacity_list;
    g.potential_list = this->potential_list;
    return g;
  }

  // b->eの辺を挿入 頂点番号は0オリジン
  void insert_edge(int b, int e, int cost, int capacity) {
    vertices_list[b].insert(e);
    cost_list[b][e] = cost;
    capacity_list[b][e] = capacity;
  }
  void delete_edge(int b, int e) {
    // b-eな辺を削除
    vertices_list[b].erase(e);
    cost_list[b].erase(e);
    capacity_list[b].erase(e);
  }

  // 頂点番号aの次数を返す
  int degree_of_vertex(int a) { return vertices_list[a].size(); }

  // a-bな辺があるか調査
  bool edge_search(int a, int b) {
    return vertices_list[a].find(b) != vertices_list[a].end();
  }

  // 頂点aとbに対してa-bとなる道が存在するかどうか調査
  bool path_search(int a, int b, set<int> visited = set<int>()) {
    visited.insert(a);
    set<int>::iterator itr;
    for (itr = vertices_list[a].begin(); itr != vertices_list[a].end(); itr++) {
      if ((*itr) == b) {
        return true;
      }
      if (visited.find(*itr) == visited.end()) { // 行ったこと無いなら
        if (path_search(*itr, b, visited)) {     // 再帰
          return true;
        }
      }
    }
    return false;
  }
  // start から goalまでの最短経路を探索:
  Shortest_path_result solve_dijkstra(int start, int goal) {
    set<int> visited = set<int>();
    vector<int> distance = vector<int>(n, INF);
    vector<int> predecessor = vector<int>(n);
    priority_queue<pair<int, int>, vector<pair<int, int>>,
                   greater<pair<int, int>>>
        pq;
    pq.push(pair<int, int>(0, start)); //<始点からの距離,頂点>
    while (!pq.empty()) {
      pair<int, int> p = pq.top();
      pq.pop();
      int nv = p.second; // 今の場所
      // いい道じゃなかったら
      if (distance[nv] < p.first) {
        continue;
      }
      distance[nv] = p.first;
      // いい道が見つかったら
      // nvから伸びる点をキューへ
      for (set<int>::iterator itr = vertices_list[nv].begin();
           itr != vertices_list[nv].end(); itr++) {
        int next = (*itr);
        if (distance[next] > distance[nv] + cost_list[nv][next]) {
          distance[next] = distance[nv] + cost_list[nv][next];
          predecessor[next] = nv;
          pq.push(pair<int, int>(distance[next], next));
        }
      }
    }
    // 答えを整形
    Shortest_path_result result;
    result.path = vector<int>();
    if (distance[goal] == INF) {
    } else {
      result.path.push_back(goal);
      if (start != goal) {
        while (true) {
          int now = result.path.back();
          int pre = predecessor[now];
          result.path.push_back(pre);
          if (pre == start) {
            reverse(result.path.begin(), result.path.end());
            break;
          }
        }
      }
    }
    result.sum_of_cost = distance[goal];
    result.distance = distance;
    result.predecessor = predecessor;
    return result;
  }
  // mcf2
  pair<int, vector<Edge>> solve_mincostflow(int s, int t, int flow_size) {
    int answer_cost = 0;
    vector<map<int, int>> origin_cost = cost_list;
    vector<map<int, int>> answer_flow =
        vector<map<int, int>>(n, map<int, int>());

    Graph residual_g = clone();

    while (flow_size > 0) {
      Shortest_path_result res;
      vector<int> path;
      int min_capa = INF;
      // 補助ネットワークから最小コストのパスを検索
      res = residual_g.solve_dijkstra(s, t);
      path = res.path;
      // tに到達不可能だったら
      if (res.distance[t] == INF) {
        // cerr << "infeasible" << endl;
        return pair<int, vector<Edge>>(-1, vector<Edge>());
      }
      // ポテンシャルを更新
      for (int i = 0; i < n; i++) {
        potential_list[i] = potential_list[i] - res.distance[i];
      }

      // 流せる量を計算
      min_capa = INF;
      for (int to = t; to != s; to = res.predecessor[to]) {
        int from = res.predecessor[to];
        min_capa = min(min_capa, residual_g.capacity_list[from][to]);
      }
      min_capa = min(min_capa, flow_size);
      // flow_sizeをへらす
      flow_size -= min_capa;

      // フローを記録
      for (int to = t; to != s; to = res.predecessor[to]) {
        int from = res.predecessor[to];
        if (vertices_list[from].find(to) != vertices_list[from].end()) {
          if (answer_flow[from].find(to) !=
              answer_flow[from].end()) { // from -> toがあったら
            answer_flow[from][to] += min_capa;
          } else { // 初流し
            answer_flow[from][to] = min_capa;
          }
        } else { // 逆向きに流れる = 減らす
          answer_flow[to][from] -= min_capa;
          if (answer_flow[to][from] == 0) {
            answer_flow[to].erase(from);
          }
        }
        // 残余ネットワークの作成
        int nflw = min_capa; // from->toにnflwだけ流れてる
        int old_capa = residual_g.capacity_list[from][to];
        // 元の辺を削除
        residual_g.delete_edge(from, to);
        // to -> fromにnflwだけながす
        if (residual_g.vertices_list[to].find(from) !=
            residual_g.vertices_list[to].end()) { // 流した痕跡あり
          residual_g.capacity_list[to][from] += nflw;
        } else { // 流した痕跡なし(新たに逆辺を追加
          residual_g.insert_edge(to, from, 0, nflw);
        }
        // from->toの辺の容量を減らす
        if (old_capa != nflw) { // あまりあり
          residual_g.insert_edge(from, to, 0, old_capa - nflw);
        }
      }
      // 辺のコストの更新
      // 残余ネットワークの全ての辺のコストを更新
      set<int>::iterator itr;
      for (int i = 0; i < n; i++) {
        for (itr = residual_g.vertices_list[i].begin();
             itr != residual_g.vertices_list[i].end(); itr++) {
          int from = i;
          int to = *itr;
          if (origin_cost[from].find(to) != origin_cost[from].end()) {
            residual_g.cost_list[from][to] = origin_cost[from][to] -
                                             potential_list[from] +
                                             potential_list[to];
          } else {
            residual_g.cost_list[from][to] = -1 * origin_cost[to][from] -
                                             potential_list[from] +
                                             potential_list[to];
          }
        }
      }
    }
    // print();
    // 解を整形

    vector<Edge> return_flow = vector<Edge>();
    answer_cost = 0;
    for (int i = 0; i < n; i++) {
      map<int, int>::iterator itr;
      for (itr = answer_flow[i].begin(); itr != answer_flow[i].end(); itr++) {
        int from = i;
        int to = itr->first;
        if (origin_cost[from].find(to) ==
            origin_cost[from].end()) { // 元々はない辺=逆辺
          /// 元々存在する辺の流量をへらす
          answer_flow[to][from] -= itr->second;
        }
      }
    }
    for (int i = 0; i < n; i++) {
      map<int, int>::iterator itr;
      for (itr = answer_flow[i].begin(); itr != answer_flow[i].end(); itr++) {
        int from = i;
        int to = itr->first;
        if (origin_cost[from].find(to) !=
            origin_cost[from].end()) { // from->toは元々ある辺
          answer_cost += itr->second * origin_cost[from][to];
          return_flow.push_back(Edge(from, to, itr->second));
        }
      }
    }

    return pair<int, vector<Edge>>(answer_cost, return_flow);
  }

  // 出力
  void print(void) {
    int i = 0;
    vector<set<int>>::iterator itr;
    set<int>::iterator itr_c;
    for (itr = vertices_list.begin(); itr != vertices_list.end(); itr++) {
      cout << i << ":";
      for (itr_c = (*itr).begin(); itr_c != (*itr).end(); itr_c++) {
        cout << *itr_c << "(" << capacity_list[i][*itr_c] << ")"
             << ",";
      }
      i++;
      cout << endl;
    }
  }
};

int main() {
  int v, e, f;
  std::cin >> v >> e >> f;
  Graph g = Graph(v);
  for (int i = 0; i < e; i++) {
    int from, to, capa, cost;
    cin >> from >> to >> capa >> cost;
    g.insert_edge(from, to, cost, capa);
  }
  std::cout << g.solve_mincostflow(0, v - 1, f).first << endl;

  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <iostream>
#include <limits>
#include <map>
#include <queue>
#include <set>
#include <stdio.h>
#include <time.h>
#include <utility>
#include <vector>

using namespace std;

class Edge {
public:
  int from;
  int to;
  int flow;
  Edge(int f, int t, int d) {
    from = f;
    to = t;
    flow = d;
  }
};

class Shortest_path_result {
public:
  int sum_of_cost;
  vector<int> path;
  vector<int> distance;
  vector<int> predecessor;
  Shortest_path_result() {}
};

class Graph {
public:
  int INF;
  int n;
  vector<set<int>> vertices_list;
  vector<map<int, int>> cost_list;
  vector<map<int, int>> capacity_list;
  vector<int> potential_list;

  Graph() {}
  Graph(int n) { // nは頂点数
    INF = 1e9;
    this->n = n;
    // 各種リストを集合
    // n個の空集合を挿入
    vertices_list.insert(vertices_list.begin(), n, set<int>());
    cost_list.insert(cost_list.begin(), n, map<int, int>());
    capacity_list.insert(capacity_list.begin(), n, map<int, int>());
    potential_list = vector<int>(n, 0);
  }

  // b->eの辺を挿入 頂点番号は0オリジン
  void insert_edge(int b, int e, int cost, int capacity) {
    vertices_list[b].insert(e);
    cost_list[b][e] = cost;
    capacity_list[b][e] = capacity;
  }
  void delete_edge(int b, int e) {
    // b-eな辺を削除
    vertices_list[b].erase(e);
    cost_list[b].erase(e);
    capacity_list[b].erase(e);
  }

  // 頂点番号aの次数を返す
  int degree_of_vertex(int a) { return vertices_list[a].size(); }

  // a-bな辺があるか調査
  bool edge_search(int a, int b) {
    return vertices_list[a].find(b) != vertices_list[a].end();
  }

  // 頂点aとbに対してa-bとなる道が存在するかどうか調査
  bool path_search(int a, int b, set<int> visited = set<int>()) {
    visited.insert(a);
    set<int>::iterator itr;
    for (itr = vertices_list[a].begin(); itr != vertices_list[a].end(); itr++) {
      if ((*itr) == b) {
        return true;
      }
      if (visited.find(*itr) == visited.end()) { // 行ったこと無いなら
        if (path_search(*itr, b, visited)) {     // 再帰
          return true;
        }
      }
    }
    return false;
  }
  // start から goalまでの最短経路を探索:
  Shortest_path_result solve_dijkstra(int start, int goal) {
    set<int> visited = set<int>();
    vector<int> distance = vector<int>(n, INF);
    vector<int> predecessor = vector<int>(n);
    priority_queue<pair<int, int>, vector<pair<int, int>>,
                   greater<pair<int, int>>>
        pq;
    pq.push(pair<int, int>(0, start)); //<始点からの距離,頂点>
    while (!pq.empty()) {
      pair<int, int> p = pq.top();
      pq.pop();
      int nv = p.second; // 今の場所
      // いい道じゃなかったら
      if (distance[nv] < p.first) {
        continue;
      }
      distance[nv] = p.first;
      // いい道が見つかったら
      for (set<int>::iterator itr = vertices_list[nv].begin();
           itr != vertices_list[nv].end(); itr++) {
        int next = (*itr);
        if (distance[next] > distance[nv] + cost_list[nv][next]) {
          distance[next] = distance[nv] + cost_list[nv][next];
          predecessor[next] = nv;
          pq.push(pair<int, int>(distance[next], next));
        }
      }
    }
    // 答えを整形
    Shortest_path_result result;

    result.path = vector<int>();
    result.path.push_back(goal);
    // 到達不可能だったら
    // if(distance[goal] == INF){
    //   return pair<int,vector<int> >(INF,vector<int>(1,start));
    // }
    // start == goalの場合
    // if(start == goal){
    //   return pair<int,vector<int> >(0,vector<int>(1,start));
    // }
    while (true) {
      int now = result.path.back();
      int pre = predecessor[now];
      result.path.push_back(pre);
      if (pre == start) {
        reverse(result.path.begin(), result.path.end());
        break;
      }
    }
    result.sum_of_cost = distance[goal];
    result.distance = distance;
    result.predecessor = predecessor;
    return result;
  }
  // mincostflow
  // sからtへflow_size流すときの最小費用流を出す
  pair<int, vector<Edge>> solve_mincostflow(int s, int t, int flow_size) {
    vector<map<int, int>> flow_list = vector<map<int, int>>(n);
    vector<map<int, int>> origin_cost = cost_list;
    int sum_flow_cost = 0;
    while (flow_size > 0) {
      Shortest_path_result res;
      vector<int> path;
      int min_capa = INF;
      res = solve_dijkstra(s, t);
      path = res.path;
      // ポテンシャルの更新
      for (int i = 0; i < n; i++) {
        potential_list[i] = potential_list[i] - res.distance[i];
      }
      // 流せる流量とコストを計算
      vector<int>::iterator itr = path.begin();
      itr++;
      for (; itr != path.end(); itr++) {
        if (min_capa > capacity_list[*(itr - 1)][*itr]) {
          min_capa = capacity_list[*(itr - 1)][*itr];
        }
      }
      if (min_capa > flow_size) {
        min_capa = flow_size;
      }
      itr = path.begin();
      itr++;
      for (; itr != path.end(); itr++) {
        if (flow_list[*(itr - 1)].find(*itr) ==
            flow_list[*(itr - 1)].end()) { // keyがないとき　
          flow_list[*(itr - 1)][*itr] = min_capa;
        } else {
          flow_list[*(itr - 1)][*itr] += min_capa;
        }
      }
      // flow_sizeを減少
      flow_size = flow_size - min_capa;
      // 辺の張替え
      itr = path.begin();
      for (itr++; itr != path.end(); itr++) {
        int capa, cost;
        int from, to;
        from = *(itr - 1);
        to = *itr;
        capa = capacity_list[from][to];
        cost = cost_list[from][to];
        delete_edge(from, to);
        if (capa - min_capa > 0) {
          insert_edge(from, to, cost,
                      capa - min_capa); // capaを減らした辺を追加　
        }
        insert_edge(to, from, -1 * cost,
                    min_capa); // 反対向きに移動した容量分の辺を追加
      }

      // 辺のコストを更新
      for (int b = 0; b > n; b++) {
        map<int, int>::iterator itr;
        for (itr = cost_list[b].begin(); itr != cost_list[b].end(); itr++) {
          (*itr).second =
              (*itr).second - potential_list[itr->first] + potential_list[b];
        }
      }
    }

    // 最終的なコストの計算
    for (int i = 0; i < n; i++) {
      map<int, int>::iterator itr;
      for (itr = flow_list[i].begin(); itr != flow_list[i].end(); itr++) {
        sum_flow_cost += origin_cost[i][itr->first] * itr->second;
        // cout << i << ","<< itr->first << "(" << itr->second << ")" << endl;
      }
    }
    cout << sum_flow_cost << endl;

    return pair<int, vector<Edge>>();
  }

  // 出力
  void print(void) {
    int i = 0;
    vector<set<int>>::iterator itr;
    set<int>::iterator itr_c;
    for (itr = vertices_list.begin(); itr != vertices_list.end(); itr++) {
      cout << i << ":";
      for (itr_c = (*itr).begin(); itr_c != (*itr).end(); itr_c++) {
        cout << *itr_c << "(" << capacity_list[i][*itr_c] << ")"
             << ",";
      }
      i++;
      cout << endl;
    }
  }
};

int main() {
  const int inf = 1e9;
  int v, e, flow; // m 財の数,  n 買い手の数
  Graph g;

  // 入力
  cin >> v >> e >> flow;
  g = Graph(v);

  int from, to, cost, cap;
  for (int i = 0; i < e; i++) {
    cin >> from >> to >> cap >> cost;
    g.insert_edge(from, to, cost, cap);
  }

  g.solve_mincostflow(0, v - 1, flow);

  return 0;
}