2023 rcpc

2023/rcpc/d1/editorial/C. Almost Tree Cut - alungu.cpp

@@ -0,0 +1,225 @@
+/// Lungu Alexandru
+#include <vector>
+#include <set>
+#include <map>
+#include <iostream>
+#include <cassert>
+#include <cmath>
+
+using namespace std;
+void debug_out() { cerr << endl; }
+template<class T> ostream& prnt(ostream& out, T v) { out << v.size() << '\n'; for(auto e : v) out << e << ' '; return out;}
+template<class T> ostream& operator<<(ostream& out, vector <T> v) { return prnt(out, v); }
+template<class T> ostream& operator<<(ostream& out, set <T> v) { return prnt(out, v); }
+template<class T1, class T2> ostream& operator<<(ostream& out, map <T1, T2> v) { return prnt(out, v); }
+template<class T1, class T2> ostream& operator<<(ostream& out, pair<T1, T2> p) { return out << '(' << p.first << ' ' << p.second << ')'; }
+template <typename Head, typename... Tail> void debug_out(Head H, Tail... T) { cerr << " " << H; debug_out(T...);}
+#define dbg(...) cerr << #__VA_ARGS__ << " ->", debug_out(__VA_ARGS__)
+#define dbg_v(x, n) do{cerr<<#x"[]: ";for(int _=0;_<n;++_)cerr<<x[_]<<" ";cerr<<'\n';}while(0)
+#define dbg_ok cerr<<"OK!\n"
+#define ll long long
+#define ld long double
+#define ull unsigned long long
+#define pii pair<int,int>
+#define MOD 1000000007
+#define zeros(x) x&(x-1)^x
+#define fi first
+#define se second
+
+const long double PI = acos(-1);
+const bool ASSERT = true;
+const int NMAX = 2 * 1e5;
+
+int n; 
+ll totalCost, dp[2 * NMAX + 5], cost[2 * NMAX + 5];
+bool vis[2 * NMAX + 5];
+vector<int> v[2 * NMAX + 5];
+
+void check(ll partA, ll partB, ll& answer)
+{
+   // dbg(partA, partB, answer);
+   if (ASSERT)
+   {
+      assert(partA > 0 && "Part A is 0");
+      assert(partB > 0 && "Part A is 0");
+      assert(partA + partB == totalCost && "Parts are invalid");
+   }
+   answer = min(answer, abs(partA - partB));
+}
+
+int dfs(int x, vector<int>& path, vector<int>& cycle, int prev = -1)
+{
+   path.push_back(x);
+   vis[x] = 1;
+   for (int next : v[x])
+   {
+      if (next == prev) continue;
+      if (vis[next])
+      {
+         while (path.back() != next)
+         {
+            cycle.push_back(path.back());
+            path.pop_back();
+         }
+         cycle.push_back(next);
+         return 1;
+      }
+      if (dfs(next, path, cycle, x)) return 1;
+   }
+   path.pop_back();
+   return 0;
+}
+
+ll solve1_subtree(int root, int prev, ll& answer)
+{
+   if (ASSERT)
+   {
+      assert(prev != root && "Invalid parameters solve1 (2)");
+      // assert(find(v[root].begin(), v[root].end(), prev) != v[root].end() && "Invalid cycle invariant in solve1 (3)");
+   }
+
+   ll subTreeCost = cost[root];
+   for (int next : v[root])
+   {
+      if (next == prev) continue;
+      subTreeCost += solve1_subtree(next, root, answer);
+   }
+
+   check(subTreeCost, totalCost - subTreeCost, answer);
+   dp[root] = subTreeCost;
+   return subTreeCost;
+}
+
+ll solve1(int x, int prev, int next)
+{
+   if (ASSERT)
+   {
+      assert(prev != next && prev != x && x != next && "Invalid parameters solve1 (1)");
+      // assert(find(v[x].begin(), v[x].end(), prev) != v[x].end() && "Invalid cycle invariant in solve1 (1)");
+      // assert(find(v[x].begin(), v[x].end(), next) != v[x].end() && "Invalid cycle invariant in solve1 (2)");
+   }
+
+   dp[x] = cost[x];
+   if (v[x].size() == 2)
+   {
+      return 1e18;
+   }
+
+   ll answer = 1e18;
+   for (int root : v[x])
+   {
+      if (root == prev || root == next) continue;
+      ll subTreeCost = solve1_subtree(root, x, answer);
+      dp[x] += subTreeCost;
+   }
+
+   if (ASSERT)
+   {
+      assert(answer < 1e18 && "Invalid answer after solve1");
+   }
+   return answer;
+}
+
+int main()
+{
+   ios::sync_with_stdio(false);
+
+   cin >> n;
+
+   ll answer = 1e18;
+
+   for (int i = 0; i < n; i++)
+   {
+      cin >> cost[i + 1];
+      totalCost += cost[i + 1];
+   }
+
+   for (int i = 0; i < n; i++)
+   {
+      int x, y;
+      cin >> x >> y;
+      v[x].push_back(y);
+      v[y].push_back(x);
+   }
+
+   vector<int> cycle;
+   vector<int> path;
+   dfs(1, path, cycle);
+
+   if (ASSERT)
+   {
+      assert(cycle.size() >= 3 && "Cycle size is too small");
+      for (int i = 0; i < cycle.size(); i++)
+      {
+         int now = cycle[i];
+         int next = i + 1 == cycle.size() ? cycle[0] : cycle[i + 1];
+         // assert(find(v[now].begin(), v[now].end(), next) != v[now].end() && "Cycle is invalid.");
+      }
+   }
+
+   // case 1: edge from cycle + edge from outer-ring
+   for (int i = 0; i < cycle.size(); i++)
+   {
+      int prev = i == 0 ? cycle[cycle.size() - 1] : cycle[i - 1];
+      int now = cycle[i];
+      int next = i + 1 == cycle.size() ? cycle[0] : cycle[i + 1];
+      answer = min(answer, solve1(now, prev, next));
+   }
+
+   // case 2: both edges from cycle
+   ll partA = dp[cycle[0]];
+   ll partB = totalCost - dp[cycle[0]];
+   int cycleSize = cycle.size();
+   for (int i = 0; i < cycleSize; i++)
+   {
+      cycle.push_back(cycle[i]);
+   }
+
+   for (int i = 0, j = 1; i < cycleSize;)
+   {
+      if (ASSERT)
+      {
+         assert(i < j && "Invalid iterators (a1)");
+         assert(j - i < cycleSize && "Invalid iterators (a2)");
+         assert(partA + partB == totalCost && "Invalid parts (a3)");
+      }
+
+      check(partA, partB, answer);
+
+      while (j - i < cycleSize - 1 && partA < partB)
+      {
+         partA += dp[cycle[j]];
+         partB -= dp[cycle[j]];
+         j++;
+         
+         check(partA, partB, answer);
+      }
+
+      if (ASSERT)
+      {
+         assert(i < j && "Invalid iterators (b1)");
+         assert(j - i < cycleSize && "Invalid iterators (b2)");
+         assert(partA + partB == totalCost && "Invalid parts (b3)");
+      }
+
+      partA -= dp[cycle[i]];
+      partB += dp[cycle[i]];
+      i++;
+
+      if (i == j)
+      {
+         partA += dp[cycle[j]];
+         partB -= dp[cycle[j]];
+         j++;
+      }
+   }
+
+   if (ASSERT)
+   {
+      assert(answer != 1e18 && "Invalid answer");
+   }
+
+   cout << answer << '\n';
+
+   return 0;
+}