Problem p01504 (C++) — buggy submission, verdict: Memory Limit Exceeded


#include <algorithm>
#include <cstring>
#include <iostream>
#include <queue>
#include <vector>

using namespace std;

typedef pair<int, int> pii;

const int MAX_V = 10001;
const int INF = 1000000000;

class UnionFindTree {
public:
  // indexの数の親ノード(ある集合の親を見つけたいときはここを直接参照せず、findを使う)
  vector<int> par;
  // indexを根とする木のランク
  vector<int> rank;
  // 木の最大値
  int treeSize;
  vector<bool> isBase;
  UnionFindTree(int initTreeSize = 1000) {
    // 引数で与えられた値を最大格納数とするUnionFindTreeの作成
    treeSize = initTreeSize;
    init();
  }
  void init() {
    for (int i = 0; i < treeSize; i++) {
      par.push_back(i);
      rank.push_back(0);
      isBase.push_back(false);
    }
  }
  // 与えられた数が格納されている木のルートを探索
  int find(int x) {
    if (par[x] == x)
      return x;
    else
      return par[x] = find(par[x]);
  }
  // 和集合をとる。ただしここでは各集合の親の付け替えは起こらない
  void unite(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y)
      return;
    if (rank[x] < rank[y]) {
      par[x] = y;
      if (isBase[find(x)] || isBase[find(y)])
        isBase[find(x)] = true;
    } else {
      par[y] = x;
      if (isBase[find(x)] || isBase[find(y)])
        isBase[find(x)] = true;
      if (rank[y] == rank[x])
        rank[x]++;
    }
  }
  bool same(int x, int y) { return find(x) == find(y); }
};

struct edge {
  int to;
  int cost;
  int idx;
};

vector<edge> G[MAX_V];
vector<int> bases;
bool used[MAX_V];
int n, t, k;
void prim(int V) {
  priority_queue<int> cand;
  UnionFindTree uft(V + 1);
  for (int i = 0; i < MAX_V; i++)
    used[i] = false;
  for (int i = 0; i < bases.size(); i++)
    uft.isBase[bases[i]] = true;
  priority_queue<pair<pii, pii>> pq;
  for (int i = 0; i < bases.size(); i++) {
    // 次の場所をコストとともに登録
    for (int j = 0; j < G[bases[i]].size(); j++) {
      edge &e = G[bases[i]][j];
      if (used[e.idx])
        continue;
      used[e.idx] = true;
      pq.push(make_pair(pii(e.cost, e.idx), pii(e.to, bases[i])));
    }
  }
  while (pq.size()) {
    pair<pii, pii> p = pq.top();
    pq.pop();
    int cost = p.first.first;
    int idx = p.first.second;
    int node = p.second.first;
    int from = p.second.second;
    // 今回のエッジを採用
    if (uft.isBase[uft.find(from)] && uft.isBase[uft.find(node)])
      cand.push(-cost);
    uft.unite(from, node);
    // 今回登録したノードからいける場所を登録
    for (int i = 0; i < G[node].size(); i++) {
      edge &e = G[node][i];
      if (used[e.idx])
        continue;
      used[e.idx] = true;
      pq.push(make_pair(pii(e.cost, e.idx), pii(e.to, node)));
    }
  }
  int res = 0;
  for (int i = 0; i < k; i++) {
    res += -cand.top();
    cand.pop();
  }
  cout << res << endl;
  cout << flush;
}

int main() {
  int cnt = 1;
  while (cin >> n >> t >> k && (n | t | k)) {
    for (int i = 0; i < MAX_V; i++)
      G[i].clear();
    for (int i = 0; i < n - 1; i++) {
      edge e;
      int a, b, c;
      cin >> a >> b >> c;
      e.to = a;
      e.cost = c;
      e.idx = i;
      G[b].push_back(e);
      e.to = b;
      G[a].push_back(e);
    }
    bases.clear();
    for (int i = 0; i < t; i++) {
      int a;
      cin >> a;
      bases.push_back(a);
    }
    cout << "Case " << cnt << ": ";
    prim(n);
    cnt++;
  }

  return 0;
}

Error: MLE

Fixed submission:


#include <algorithm>
#include <cstring>
#include <iostream>
#include <queue>
#include <vector>

using namespace std;

typedef pair<int, int> pii;

const int MAX_V = 10001;
const int INF = 1000000000;

class UnionFindTree {
public:
  // indexの数の親ノード(ある集合の親を見つけたいときはここを直接参照せず、findを使う)
  vector<int> par;
  // indexを根とする木のランク
  vector<int> rank;
  // 木の最大値
  int treeSize;
  vector<bool> isBase;
  UnionFindTree(int initTreeSize = 1000) {
    // 引数で与えられた値を最大格納数とするUnionFindTreeの作成
    treeSize = initTreeSize;
    init();
  }
  void init() {
    for (int i = 0; i < treeSize; i++) {
      par.push_back(i);
      rank.push_back(0);
      isBase.push_back(false);
    }
  }
  // 与えられた数が格納されている木のルートを探索
  int find(int x) {
    if (par[x] == x)
      return x;
    else
      return par[x] = find(par[x]);
  }
  // 和集合をとる。ただしここでは各集合の親の付け替えは起こらない
  void unite(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y)
      return;
    if (rank[x] < rank[y]) {
      par[x] = y;
      if (isBase[find(x)] || isBase[find(y)])
        isBase[find(x)] = true;
    } else {
      par[y] = x;
      if (isBase[find(x)] || isBase[find(y)])
        isBase[find(x)] = true;
      if (rank[y] == rank[x])
        rank[x]++;
    }
  }
  bool same(int x, int y) { return find(x) == find(y); }
};

struct edge {
  int to;
  int cost;
  int idx;
};

vector<edge> G[MAX_V];
vector<int> bases;
bool used[MAX_V];
int n, t, k;
void prim(int V) {
  priority_queue<int> cand;
  UnionFindTree uft(V + 1);
  for (int i = 0; i < MAX_V; i++)
    used[i] = false;
  for (int i = 0; i < bases.size(); i++)
    uft.isBase[bases[i]] = true;
  priority_queue<pair<pii, pii>> pq;
  for (int i = 0; i < bases.size(); i++) {
    // 次の場所をコストとともに登録
    for (int j = 0; j < G[bases[i]].size(); j++) {
      edge &e = G[bases[i]][j];
      pq.push(make_pair(pii(e.cost, e.idx), pii(e.to, bases[i])));
    }
  }
  while (pq.size()) {
    pair<pii, pii> p = pq.top();
    pq.pop();
    int cost = p.first.first;
    int idx = p.first.second;
    int node = p.second.first;
    int from = p.second.second;
    if (used[idx])
      continue;
    used[idx] = true;
    // 今回のエッジを採用
    if (uft.isBase[uft.find(from)] && uft.isBase[uft.find(node)])
      cand.push(-cost);
    uft.unite(from, node);
    // 今回登録したノードからいける場所を登録
    for (int i = 0; i < G[node].size(); i++) {
      edge &e = G[node][i];
      pq.push(make_pair(pii(e.cost, e.idx), pii(e.to, node)));
    }
  }
  int res = 0;
  for (int i = 0; i < k; i++) {
    res += -cand.top();
    cand.pop();
  }
  cout << res << endl;
  cout << flush;
}

int main() {
  int cnt = 1;
  while (cin >> n >> t >> k && (n | t | k)) {
    for (int i = 0; i < MAX_V; i++)
      G[i].clear();
    for (int i = 0; i < n - 1; i++) {
      edge e;
      int a, b, c;
      cin >> a >> b >> c;
      e.to = a;
      e.cost = c;
      e.idx = i;
      G[b].push_back(e);
      e.to = b;
      G[a].push_back(e);
    }
    bases.clear();
    for (int i = 0; i < t; i++) {
      int a;
      cin >> a;
      bases.push_back(a);
    }
    cout << "Case " << cnt << ": ";
    prim(n);
    cnt++;
  }

  return 0;
}