Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #include <bits/stdc++.h>
- using namespace std;
- typedef long long ll;
- typedef long double ld;
- typedef pair <int, int> ii;
- typedef vector <int> vi;
- typedef vector <vi> vvi;
- typedef vector <ii> vii;
- typedef vector <vii> vvii;
- #define endl '\n'
- #define PB push_back
- #define MP make_pair
- #define fr first
- #define sc second
- #define OO (1000000000) // ToDo
- #define EPS (1e-9) // ToDo
- #define MOD (1000000007) // ToDo
- #define all(v) ((v).begin()),((v).end())
- #define wt(x) cout<< #x <<" = "<<"\""<< (x) <<"\""<<endl
- #define FASTIO ios::sync_with_stdio(0), cin.tie(0), cout.tie(0)
- void read_file(bool outToFile = true){
- #ifdef LOCAL_TEST
- freopen("in", "rt", stdin);
- if(outToFile)
- freopen("out", "wt", stdout);
- #endif
- }
- //
- const int MAXN = 100 * 1000 + 99;
- //
- int TC;
- int n;
- int W[MAXN];
- vvi G;
- ll dp[MAXN], depth[MAXN];
- ll sub[MAXN];
- int P[MAXN]; // parent array
- ll S;
- //
- void DFS(int u=0, int p=-1){
- P[u] = p;
- depth[u] = p == -1? 0 : depth[p] + 1;
- sub[u] = 1;
- ll X = 0, Y = 0; // sizes of the left, right subtrees
- for(int i=0; i<G[u].size(); i++)
- {
- int v = G[u][i];
- if(v == p) continue;
- DFS(v, u);
- sub[u] += sub[v];
- if(X == 0)
- X = sub[v];
- else
- Y = sub[v];
- }
- ll Z = n-1 - X - Y; // size of other nodes
- dp[u] = X*Y + Y*Z + Z*X; // paths that don't end at u
- dp[u] = dp[u] + n-1; // adding n-1 paths which have u as an endpoint
- dp[u] = dp[u]*2; // multiplying by 2 for the two directions
- dp[u] = dp[u] + 1; // adding 1 for the path (u, u)
- }
- void solve(){
- // calculations
- DFS();
- // step one
- S = 0;
- for(int u=0; u<n; u++)
- S += W[u] * dp[u];
- if(S == 0)
- {
- printf("0\n");
- return;
- }
- // step two
- for(int u=0; u<n; u++)
- {
- if(S%dp[u] == 0) // dp[u] must be >= 1, no special cases
- {
- printf("1 ");
- printf("%d\n", u+1);
- //cerr << W[u] - S / dp[u] <<endl;
- return;
- }
- }
- // step three
- int cho = 0;
- for(int u=0; u<n; u++)
- {
- if(depth[cho] < depth[u])
- cho = u;
- }
- int u = cho, v = P[u];
- printf("2 ");
- printf("%d %d\n", u+1, v+1);
- assert(dp[u] == 2*n-1 && dp[v] == 6*n-11);
- assert(__gcd(dp[u], dp[v]) == 1); // because it's free
- //cerr<< S <<endl;
- }
- //
- int main()
- {
- //read_file();
- scanf("%d", &TC);
- while(TC--)
- {
- // reading input
- scanf("%d", &n);
- G.assign(n, vi());
- for(int u=0; u<n; u++)
- scanf("%lld", &W[u]);
- for(int i=0; i<n-1; i++)
- {
- int u, v;
- scanf("%d %d", &u, &v);
- --u;--v;
- G[u].PB(v);
- G[v].PB(u);
- }
- solve();
- //printf("\n");
- }
- }
Advertisement
Add Comment
Please, Sign In to add comment