usaco

Clean implementations of solutions to USACO problems

  1. 1
  2. 2
  3. 3
  4. 4
  5. 5
  6. 6
  7. 7
  8. 8
  9. 9
  10. 10
  11. 11
  12. 12
  13. 13
  14. 14
  15. 15
  16. 16
  17. 17
  18. 18
  19. 19
  20. 20
  21. 21
  22. 22
  23. 23
  24. 24
  25. 25
  26. 26
  27. 27
  28. 28
  29. 29
  30. 30
  31. 31
  32. 32
  33. 33
  34. 34
  35. 35
  36. 36
  37. 37
  38. 38
  39. 39
  40. 40
  41. 41
  42. 42
  43. 43
  44. 44
  45. 45
  46. 46
  47. 47
  48. 48
  49. 49
  50. 50
  51. 51
  52. 52
  53. 53
  54. 54
  55. 55
  56. 56
  57. 57
  58. 58
  59. 59
  60. 60
  61. 61
  62. 62
  63. 63
  64. 64
  65. 65
  66. 66
  67. 67
  68. 68
  69. 69
  70. 70
  71. 71
  72. 72
  73. 73
  74. 74
  75. 75
  76. 76
  77. 77
  78. 78
  79. 79
  80. 80
  81. 81
  82. 82
  83. 83
  84. 84
  85. 85
  86. 86
  87. 87
  88. 88
  89. 89
  90. 90
  91. 91
  92. 92
  93. 93
  94. 94
  95. 95
  96. 96
  97. 97
#include <algorithm>
#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include <queue>
#include <set>
#include <map>
#include <unordered_set>
#include <unordered_map>
#include <cmath>
#include <cstring>
#define init_io(pname) ifstream cin((string)pname+".in"); ofstream cout((string)pname+".out"); ios_base::sync_with_stdio(false); cin.tie(NULL)
#define FOR(i, A, B, in) for (int i = (A); i < (B); i += in)
#define REP(i, A, B) for (int i = (A); i < (B); i++)
#define RFOR(i, A, B, in) for (int i = (A) - 1; i >= (B); i -= in)
#define RREP(i, A, B) for (int i = (A) - 1; i >= (B); i--)
#define trav(A, x) for (auto& A : x)
#define mp make_pair
#define pb push_back
#define lb lower_bound
#define ub upper_bound
#define sz(x) (int)x.size()
#define all(x) begin(x), end(x)
#define rsz resize
#define mem(A, B) memset(A, (B), sizeof(A))
#define uset unordered_set
#define umap unordered_map
using namespace std;
typedef string str;
typedef long long ll;
typedef long double ld;
typedef pair<int, int> ii; typedef pair<ll, ll> pl; typedef pair<ld, ld> pd;
typedef vector<int> vi; typedef vector<ll> vl; typedef vector<ld> vd;
typedef vector<ii> vii; typedef vector<pl> vpl; typedef vector<pd> vpd;
constexpr auto INF = (int)1e9;
constexpr auto LINF = (ll)1e18;

vii G[100005];
umap<int, int> A[100005], B[100005];

ll dfs(int u, int p, int s) {
	ll ret = 0;
	for (auto& v : G[u]) { // Traverse children
		if (v.first != p) ret += dfs(v.first, u, s + (v.second ? 1 : -1));
	}
	int m = 0; // Computed largest set to merge small to large
	for (auto& v : G[u]) {
		if (v.first != p && A[v.first].size() + B[v.first].size() > A[m].size() + B[m].size()) m = v.first;
	}
	if (A[m].find(s) != A[m].end()) {
		ret += B[m][s];
		B[u][s] += A[m][s];
		A[m].erase(s);
	}
	if (A[u].size() < A[m].size()) swap(A[u], A[m]);
	if (B[u].size() < B[m].size()) swap(B[u], B[m]);
	for (auto& x : A[m]) A[u][x.first] += x.second;
	for (auto& x : B[m]) B[u][x.first] += x.second;
	// Merge other sets
	for (auto& v : G[u]) {
		if (v.first != p && v.first != m) {
			for (auto& x : A[v.first]) {
				if (B[u].find(2 * s - x.first) != B[u].end()) ret += (ll)x.second * B[u][2 * s - x.first];
			}
			for (auto& x : B[v.first]) {
				if (A[u].find(2 * s - x.first) != A[u].end()) ret += (ll)x.second * A[u][2 * s - x.first];
				if (B[u].find(2 * s - x.first) != B[u].end()) ret += (ll)x.second * B[u][2 * s - x.first];
			}
			if (A[v.first].find(s) != A[v.first].end()) {
				ret += B[v.first][s];
				B[u][s] += A[v.first][s];
				A[v.first].erase(s);
			}
			if (A[u].size() < A[v.first].size()) swap(A[u], A[v.first]);
			if (B[u].size() < B[v.first].size()) swap(B[u], B[v.first]);
			for (auto& x : A[v.first]) A[u][x.first] += x.second; // Merge "A" sets
			for (auto& x : B[v.first]) B[u][x.first] += x.second; // Merge "B" sets
		}
	}
	A[u][s]++;
	return ret;
}

int main() {
	init_io("yinyang");

	int N;
	cin >> N;
	for (int i = 0; i < N - 1; i++) {
		int A, B, t;
		cin >> A >> B >> t;
		G[A].emplace_back(B, t);
		G[B].emplace_back(A, t);
	}
	cout << dfs(1, 0, 0) << endl;
}