fork(5) download
copy 
  1. #include<bits/stdc++.h>
  2. #include <ext/pb_ds/assoc_container.hpp>
  3. #include <ext/pb_ds/tree_policy.hpp>
  4. using namespace std;
  5. using namespace __gnu_pbds;
  6.  
  7. typedef long long ll;
  8. typedef long double ld;
  9. typedef pair<int, int> pii;
  10. typedef pair<ll, ll> pll;
  11. typedef vector<int> vi;
  12. typedef vector<ll> vll;
  13.  
  14. #define ordered_set tree<int, null_type, less<int>, rb_tree_tag, tree_order_statistics_node_update>
  15. #define ordered_multiset tree<int, null_type, less_equal<int>, rb_tree_tag, tree_order_statistics_node_update>
  16. #define pb push_back
  17. #define ff first
  18. #define ss second
  19.  
  20. const int N = 3e5 + 5;
  21.  
  22. vi graf[N], sciezka, trasa;
  23. vll ile;
  24.  
  25. ll A[N], B[N];
  26. int c[N], oryg[N], czy[N], dp1[N], dp2[N];
  27.  
  28. int best, start;
  29.  
  30. bool cmp(int a, int b){
  31.     return c[a] > c[b];
  32. }
  33.  
  34. void dfs(int v, int ojc){
  35.     c[v] = oryg[v];
  36.     czy[v] = (c[v] == 0);
  37.     dp1[v] = 0;
  38.     dp2[v] = -1;
  39.  
  40.     for(int u : graf[v]){
  41.         if(u == ojc) continue;
  42.  
  43.         dfs(u, v);
  44.  
  45.         c[v] += c[u];
  46.         czy[v] &= czy[u];
  47.  
  48.         int kand = dp1[u] + 1 - 2 * (c[u] > 0);
  49.  
  50.         if(!czy[u] && kand >= dp1[v]){
  51.             dp1[v] = kand;
  52.             dp2[v] = u;
  53.         }
  54.     }
  55.  
  56.     for(int i = 0; i < (int)graf[v].size(); i++){
  57.         int u = graf[v][i];
  58.  
  59.         if(u == ojc) continue;
  60.  
  61.         if(czy[u]){
  62.             swap(graf[v][i], graf[v].back());
  63.             graf[v].pop_back();
  64.             i--;
  65.         }
  66.     }
  67. }
  68.  
  69. void dfs2(int v, int ojc, int up){
  70.     if(c[v] < 0) up -= 2;
  71.  
  72.     int kand = max(dp1[v], up);
  73.  
  74.     if(kand > best){
  75.         best = kand;
  76.         start = v;
  77.     }
  78.  
  79.     int m = graf[v].size();
  80.  
  81.     vi pref(m + 1, 0), suf(m + 1, 0);
  82.  
  83.     for(int i = 0; i < m; i++){
  84.         pref[i + 1] = pref[i];
  85.  
  86.         int u = graf[v][i];
  87.  
  88.         if(u == ojc) continue;
  89.  
  90.         int kand2 = dp1[u] + 1 - 2 * (c[u] > 0);
  91.  
  92.         pref[i + 1] = max(pref[i + 1], kand2);
  93.     }
  94.  
  95.     for(int i = m - 1; i >= 0; i--){
  96.         suf[i] = suf[i + 1];
  97.  
  98.         int u = graf[v][i];
  99.  
  100.         if(u == ojc) continue;
  101.  
  102.         int kand2 = dp1[u] + 1 - 2 * (c[u] > 0);
  103.  
  104.         suf[i] = max(suf[i], kand2);
  105.     }
  106.  
  107.     for(int i = 0; i < m; i++){
  108.         int u = graf[v][i];
  109.  
  110.         if(u == ojc) continue;
  111.  
  112.         int nast = max({pref[i], suf[i + 1], up}) + 1;
  113.  
  114.         dfs2(u, v, nast);
  115.     }
  116. }
  117.  
  118. void odz(int v){
  119.     sciezka.pb(v);
  120.  
  121.     if(dp2[v] == -1) return;
  122.  
  123.     odz(dp2[v]);
  124. }
  125.  
  126. void dfs3(int v, int ojc){
  127.     sort(graf[v].begin(), graf[v].end(), cmp);
  128.  
  129.     trasa.pb(v);
  130.     ile.pb(-A[v]);
  131.  
  132.     for(int u : graf[v]){
  133.         if(u == ojc) continue;
  134.  
  135.         dfs3(u, v);
  136.  
  137.         trasa.pb(v);
  138.         ile.pb(0);
  139.     }
  140.  
  141.     ile.back() += B[v];
  142. }
  143.  
  144. void solve(){
  145.     int n;
  146.     cin >> n;
  147.  
  148.     best = -1;
  149.     start = -1;
  150.  
  151.     sciezka.clear();
  152.     trasa.clear();
  153.     ile.clear();
  154.  
  155.     for(int i = 1; i <= n; i++){
  156.         graf[i].clear();
  157.         czy[i] = 0;
  158.         c[i] = 0;
  159.         oryg[i] = 0;
  160.         dp1[i] = 0;
  161.         dp2[i] = -1;
  162.     }
  163.  
  164.     for(int i = 1; i <= n; i++)
  165.         cin >> A[i];
  166.  
  167.     for(int i = 1; i <= n; i++)
  168.         cin >> B[i];
  169.  
  170.     for(int i = 1; i <= n; i++){
  171.         // Trzymamy dokładnie tę samą konwencję co bike.h:
  172.         // d = A - B
  173.         oryg[i] = A[i] - B[i];
  174.         c[i] = oryg[i];
  175.     }
  176.  
  177.     for(int i = 1; i < n; i++){
  178.         int a, b;
  179.         cin >> a >> b;
  180.  
  181.         // input ma numery 0..n-1
  182.         ++a;
  183.         ++b;
  184.  
  185.         graf[a].pb(b);
  186.         graf[b].pb(a);
  187.     }
  188.  
  189.     int root = -1;
  190.  
  191.     for(int i = 1; i <= n; i++){
  192.         if(A[i] != B[i]){
  193.             root = i;
  194.             break;
  195.         }
  196.     }
  197.  
  198.     // Wszystko już jest poprawne.
  199.     if(root == -1){
  200.         cout << "0\n";
  201.         cout << "0\n";
  202.         cout << "0\n";
  203.         return;
  204.     }
  205.  
  206.     // Pierwszy DP
  207.     dfs(root, 0);
  208.  
  209.     // Szukamy najlepszego początku głównej ścieżki
  210.     dfs2(root, 0, 0);
  211.  
  212.     // Drugi DFS od znalezionego początku
  213.     // i odbudowanie c jako bilansów poddrzew
  214.     for(int i = 1; i <= n; i++){
  215.         c[i] = oryg[i];
  216.         czy[i] = 0;
  217.         dp1[i] = 0;
  218.         dp2[i] = -1;
  219.     }
  220.  
  221.     dfs(start, 0);
  222.  
  223.     // Odtwarzamy główną ścieżkę
  224.     sciezka.clear();
  225.  
  226.     for(int v = start; v != -1; v = dp2[v])
  227.         sciezka.pb(v);
  228.  
  229.     /*
  230.         Usuwamy krawędzie głównej ścieżki.
  231.         To jest dokładnie to, co robi oryginalny kod bike.h.
  232.     */
  233.     for(int i = 1; i < (int)sciezka.size(); i++){
  234.         int x = sciezka[i - 1];
  235.         int y = sciezka[i];
  236.  
  237.         auto it1 = find(graf[x].begin(), graf[x].end(), y);
  238.         graf[x].erase(it1);
  239.  
  240.         auto it2 = find(graf[y].begin(), graf[y].end(), x);
  241.         graf[y].erase(it2);
  242.     }
  243.  
  244.     /*
  245.         Budowanie konkretnej trasy.
  246.     */
  247.     int m = sciezka.size();
  248.  
  249.     for(int i = 0; i < m; i++){
  250.         if(i + 1 == m || c[sciezka[i + 1]] <= 0){
  251.  
  252.             int j = i - 1;
  253.  
  254.             while(j >= 0 && c[sciezka[j + 1]] > 0)
  255.                 j--;
  256.  
  257.             for(int k = i; k > j; k--)
  258.                 dfs3(sciezka[k], -1);
  259.  
  260.             for(int k = j + 2; k <= i; k++){
  261.                 trasa.pb(sciezka[k]);
  262.                 ile.pb(0);
  263.             }
  264.         }
  265.         else{
  266.             trasa.pb(sciezka[i]);
  267.             ile.pb(0);
  268.         }
  269.     }
  270.  
  271.     int koszt = (int)trasa.size() - 1;
  272.  
  273.     cout << koszt << '\n';
  274.  
  275.     // z powrotem na numerację 0..n-1
  276.     for(int i = 0; i < (int)trasa.size(); i++){
  277.         if(i) cout << ' ';
  278.         cout << trasa[i] - 1;
  279.     }
  280.     cout << '\n';
  281.  
  282.     for(int i = 0; i < (int)ile.size(); i++){
  283.         if(i) cout << ' ';
  284.         cout << ile[i];
  285.     }
  286.     cout << '\n';
  287. }
  288.  
  289. int main(){
  290.     ios_base::sync_with_stdio(0);
  291.     cin.tie(0);
  292.  
  293.     int t;
  294.     cin >> t;
  295.  
  296.     while(t--)
  297.         solve();
  298.  
  299.     return 0;
  300. }
Success #stdin #stdout 0.01s 17120KB
comments (?)
stdin
copy 
2
4
1 0 1 5 0
1 0 0 3 3
0 1
1 2
1 3
5
3 0 1 2 2
2 2 1 3 0
2 0
2 4
2 3
2 1
stdout
copy 
3
2 1 0 1
-1 0 -1 1
0
0
2
https://ideone.com/1NCFMN
language:
C++ (gcc 8.3)
created:
1 day ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!