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

#include "bits/stdc++.h"
#include <unordered_map>
#include <unordered_set>
#pragma warning(disable : 4996)
using namespace std;
using ld = long double;
template <class T> using Table = vector<vector<T>>;
const ld eps = 1e-9;

//// < "D:\\D_Download\\Visual Studio
///2015\\Projects\\programing_contest_c++\\Debug\\a.txt" > "D:\\D_Download\\Visual
///Studio 2015\\Projects\\programing_contest_c++\\Debug\\b.answer"
struct aa {
  int root;
  int from;
  int data;
  long long int time;
};
typedef int Weight;
map<pair<int, int>, vector<int>> longs;
struct Edge {
  int src, dst;
  Weight weight;
  Edge(int src_, int dst_, Weight weight_)
      : src(src_), dst(dst_), weight(weight_) {}
  Edge(int src_, int dst_) : src(src_), dst(dst_) {}
};
class Compare {
public:
  // aa?????????????????¶
  bool operator()(const aa &l, const aa &r) { return l.time < r.time; }
};
vector<int> datas;
vector<vector<Edge>> edges;
pair<long long int, vector<int>> dfs(const int now, const int from) {
  long long int amax = 0;
  vector<int> max_vs;
  for (auto e : edges[now]) {
    if (e.dst == from)
      continue;
    auto p(dfs(e.dst, e.src));
    if (amax < p.first + e.weight) {
      amax = p.first + e.weight;
      max_vs = p.second;
    }
  }
  max_vs.emplace_back(now);
  return make_pair(amax, max_vs);
}
priority_queue<aa, vector<aa>, Compare> que;
void dfs2(const int now, const int from, const vector<int> &max_vs, int &depth,
          priority_queue<aa, vector<aa>, Compare> &que, const int data_id) {

  for (auto e : edges[now]) {
    if (e.dst == from)
      continue;
    if (e.dst == max_vs[depth + 1]) {
      depth++;
      dfs2(e.dst, e.src, max_vs, depth, que, data_id);
      depth--;
    } else {
      auto p = dfs(e.dst, e.src);
      reverse(p.second.begin(), p.second.end());
      if (longs.find(make_pair(e.dst, e.src)) == longs.end()) {
        longs[make_pair(e.dst, e.src)] = p.second;
      }
      que.push(
          aa{e.dst, e.src, data_id, (p.first + e.weight) * datas[data_id]});
    }
  }
  return;
}
int main() {
  int N, M;
  cin >> N >> M;
  edges.resize(N);
  for (int i = 0; i < N - 1; ++i) {
    int a, b, c;
    cin >> a >> b >> c;
    a--;
    b--;
    edges[a].push_back(Edge(a, b, c));
    edges[b].push_back(Edge(b, a, c));
  }
  for (int i = 0; i < N; ++i) {
    int a;
    cin >> a;
    datas.emplace_back(a);
  }
  for (int i = 0; i < N; ++i) {
    auto p(dfs(i, -1));
    reverse(p.second.begin(), p.second.end());
    if (longs.find(make_pair(i, -1)) == longs.end()) {
      longs[make_pair(i, -1)] = p.second;
    }
    que.push(aa{i, -1, i, p.first * datas[i]});
  }
  long long int ans = 0;
  while (!que.empty()) {
    if (!M)
      break;
    M--;
    aa atop(que.top());
    que.pop();
    ans += atop.time;
    int d = 0;
    dfs2(atop.root, atop.from, longs[make_pair(atop.root, atop.from)], d, que,
         atop.data);
  }
  cout << ans << endl;
  return 0;
}

Error: MLE

Fixed submission:

#include "bits/stdc++.h"
#include <unordered_map>
#include <unordered_set>
#pragma warning(disable : 4996)
using namespace std;
using ld = long double;
template <class T> using Table = vector<vector<T>>;
const ld eps = 1e-9;

//// < "D:\\D_Download\\Visual Studio
///2015\\Projects\\programing_contest_c++\\Debug\\a.txt" > "D:\\D_Download\\Visual
///Studio 2015\\Projects\\programing_contest_c++\\Debug\\b.answer"
struct aa {
  int root;
  int from;
  vector<int> longs;
  int data;
  long long int time;
};
typedef int Weight;
struct Edge {
  int src, dst;
  Weight weight;
  Edge(int src_, int dst_, Weight weight_)
      : src(src_), dst(dst_), weight(weight_) {}
  Edge(int src_, int dst_) : src(src_), dst(dst_) {}
};
class Compare {
public:
  // aa?????????????????¶
  bool operator()(const aa &l, const aa &r) { return l.time < r.time; }
};
vector<int> datas;
vector<vector<Edge>> edges;
pair<long long int, vector<int>> dfs(const int now, const int from) {
  long long int amax = 0;
  vector<int> max_vs;
  for (auto e : edges[now]) {
    if (e.dst == from)
      continue;
    auto p(dfs(e.dst, e.src));
    if (amax < p.first + e.weight) {
      amax = p.first + e.weight;
      max_vs = p.second;
    }
  }
  max_vs.emplace_back(now);
  return make_pair(amax, max_vs);
}
priority_queue<aa, vector<aa>, Compare> que;
void dfs2(const int now, const int from, const vector<int> &max_vs, int &depth,
          priority_queue<aa, vector<aa>, Compare> &que, const int data_id) {

  for (auto e : edges[now]) {
    if (e.dst == from)
      continue;
    if (e.dst == max_vs[depth + 1]) {
      depth++;
      dfs2(e.dst, e.src, max_vs, depth, que, data_id);
      depth--;
    } else {
      auto p = dfs(e.dst, e.src);
      reverse(p.second.begin(), p.second.end());
      que.push(aa{e.dst, e.src, p.second, data_id,
                  (p.first + e.weight) * datas[data_id]});
    }
  }
  return;
}
int main() {
  int N, M;
  cin >> N >> M;
  edges.resize(N);
  for (int i = 0; i < N - 1; ++i) {
    int a, b, c;
    cin >> a >> b >> c;
    a--;
    b--;
    edges[a].push_back(Edge(a, b, c));
    edges[b].push_back(Edge(b, a, c));
  }
  for (int i = 0; i < N; ++i) {
    int a;
    cin >> a;
    datas.emplace_back(a);
  }
  for (int i = 0; i < N; ++i) {
    auto p(dfs(i, -1));
    reverse(p.second.begin(), p.second.end());
    que.push(aa{i, -1, p.second, i, p.first * datas[i]});
  }
  long long int ans = 0;
  while (!que.empty()) {
    if (!M)
      break;
    M--;
    aa atop(que.top());
    que.pop();
    ans += atop.time;
    int d = 0;
    dfs2(atop.root, atop.from, atop.longs, d, que, atop.data);
  }
  cout << ans << endl;
  return 0;
}