Cod sursa(job #1070294)

Utilizator bghimisFMI Ghimis Bogdan bghimis Data 31 decembrie 2013 16:31:21
Problema Cerere Scor 0
Compilator cpp Status done
Runda Arhiva de probleme Marime 1.93 kb
#include <stack>
#include <bitset>
#include <fstream>
#include <vector>
#include <deque>
#include <algorithm>
#include <unordered_map>
using namespace std;

ifstream in ("cerere.in");
ofstream out("cerere.out");

int v[100001];
unordered_map<int, int> raspuns;
vector<int> relatii[100001];

void DFS(int nodStart, vector<int>* relatii)
{
  bitset<100001> vizitat;

  stack<int> myStack;
  myStack.push(nodStart);
  vizitat[nodStart] = true;

  unordered_map<int, int> drum;
  while (!myStack.empty())
    {
      int nod_crt = myStack.top();
      myStack.pop();

      bool frunza = true;
      for (auto &it : relatii[nod_crt])
	if (vizitat[it] == false)
	  {
	    myStack.push(it);
	    vizitat[it] = true;
	    frunza = false;

	    drum[it] = nod_crt;
	  }
      
      if (frunza == true)
	{
	  deque <int> coada;
	  for (auto it = nod_crt; it != 0; it = drum[it])
	    coada.push_front(it);

	  vector <int> vect;
	  for (auto &it : coada)
	    vect.push_back(it);

	  int size = vect.size();
	  for (int i = 0; i < size; ++i)
	    {
	      int nod_crt = vect[i];
	      int cnt = 0;

	      if (v[nod_crt] == 0)
		{
		  raspuns[nod_crt] = 1;
		  continue;
		}
	      
	      int nou_i = i;
	      while (v[nod_crt] != 0 && raspuns[nod_crt] == 0)
		{
		  ++cnt;
		  nod_crt = vect [nou_i - v[nod_crt]];
		  nou_i = nou_i - v[nod_crt];
		}

	      if (raspuns[nod_crt] != 0)
		raspuns[vect[i]] = cnt + raspuns[nod_crt];
	      else
		raspuns[vect[i]] = cnt;
	    }
	}
    }
}

int main()
{
  int n;
  in >> n;
  
  for (int i = 1; i <= n; ++i)
    in >> v[i];

  int x, y;
  for (int i = 1; i < n; ++i)
    {
      in >> x >> y;
      relatii[x].push_back(y);
    }

  DFS(1, relatii);

  vector <pair<int, int>> r;
  for (auto &it : raspuns)
    r.push_back(it);

  sort (r.begin(), r.end());

  for (auto &it : r)
    out <<  it.second - 1 << " ";

  return 0;
}