test/aoj/dsl_1_b.test.cpp

View this file on GitHub
Last update: 2022-12-18 22:46:34+09:00
Problem: https://onlinejudge.u-aizu.ac.jp/courses/library/3/DSL/1/DSL_1_B

Depends on

src/structure/weighted_union_find_tree.hpp

Code

#define PROBLEM                                                                \
  "https://onlinejudge.u-aizu.ac.jp/courses/library/3/DSL/1/DSL_1_B"

#include <bits/stdc++.h>
using namespace std;

#include "../../src/structure/weighted_union_find_tree.hpp"

int main() {
  int n, q;
  cin >> n >> q;
  WeightedUnionFindTree<long long> uft(n);
  for (int i = 0; i < q; ++i) {
    int com;
    cin >> com;
    if (com == 0) {
      int x, y, z;
      cin >> x >> y >> z;
      uft.unite(x, y, z);
    } else {
      int x, y;
      cin >> x >> y;
      if (uft.is_same(x, y)) {
        cout << uft.diff(x, y) << endl;
      } else {
        cout << "?" << endl;
      }
    }
  }
  return 0;
}

#line 1 "test/aoj/dsl_1_b.test.cpp"
#define PROBLEM                                                                \
  "https://onlinejudge.u-aizu.ac.jp/courses/library/3/DSL/1/DSL_1_B"

#include <bits/stdc++.h>
using namespace std;

#line 2 "src/structure/weighted_union_find_tree.hpp"

#line 4 "src/structure/weighted_union_find_tree.hpp"
using namespace std;

template <typename T> struct WeightedUnionFindTree {
private:
  int n;
  vector<int> parent, rank;
  vector<T> weight;

public:
  WeightedUnionFindTree(int n)
      : n(n), parent(n, -1), rank(n, 0), weight(n, 0) {}

  int root(int x) {
    if (parent[x] == -1) {
      return x;
    }
    int p = root(parent[x]);
    weight[x] += weight[parent[x]];
    return parent[x] = p;
  }

  void unite(int x, int y, T w) {
    w += calc_weight(x);
    w -= calc_weight(y);
    x = root(x);
    y = root(y);
    if (x == y) {
      return;
    }
    if (rank[x] < rank[y]) {
      swap(x, y);
      w = -w;
    }
    if (rank[x] == rank[y]) {
      ++rank[x];
    }
    parent[y] = x;
    weight[y] = w;
  }

  bool is_same(int x, int y) { return root(x) == root(y); }

  T calc_weight(int x) {
    root(x);
    return weight[x];
  }

  T diff(int x, int y) { return calc_weight(y) - calc_weight(x); }
};
#line 8 "test/aoj/dsl_1_b.test.cpp"

int main() {
  int n, q;
  cin >> n >> q;
  WeightedUnionFindTree<long long> uft(n);
  for (int i = 0; i < q; ++i) {
    int com;
    cin >> com;
    if (com == 0) {
      int x, y, z;
      cin >> x >> y >> z;
      uft.unite(x, y, z);
    } else {
      int x, y;
      cin >> x >> y;
      if (uft.is_same(x, y)) {
        cout << uft.diff(x, y) << endl;
      } else {
        cout << "?" << endl;
      }
    }
  }
  return 0;
}