#include <numeric>
#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
vector<long> cumulativeNeighbourPigs;
void init(vector<long>& pigs, vector<vector<int>>& adj)
{
cumulativeNeighbourPigs.resize(pigs.size());
for(int i = 0; i<pigs.size(); ++i)
{
for(auto child : adj[i])
{
cumulativeNeighbourPigs[i] += pigs[child];
}
}
}
long DFS(int startNode, int prevN, vector<vector<int>>& adj, vector<long>& pig, int depth, const int v, long currP)
{
// cout<<"currNode: "<<startNode<<endl;
if(depth == v)
{
return currP;
}
else
{
if(depth > v)
return -1;
else
{
long currMax = 0;
for(auto child : adj[startNode])
{
if(child == prevN)
continue;
// cout<<"trying child: "<<child<<endl;
long childPigeons = cumulativeNeighbourPigs[child] - pig[startNode];
if(adj[child].size() > 1)
childPigeons = DFS(child, startNode,adj, pig, depth+1, v, currP+ childPigeons);
if(currMax < childPigeons)
currMax = childPigeons;
// cout<<"new max "<<currMax<<endl;
}
return max(currP, currMax);
}
}
}
long maxPigeonsPath(vector<vector<int>>& adj, vector<long>& pig, const int v)
{
if(v == 0)
{
return 0;
}
long maxP = 0;
for(int i =0 ;i<adj.size(); ++i)
{
long pigs = DFS(i,-1, adj, pig, 1, v, cumulativeNeighbourPigs[i]);
if(pigs > maxP)
maxP = pigs;
}
return maxP;
}
int main()
{
int n; cin>>n;
int v; cin>>v;
vector<long> pigeons(n);
for(int i =0; i<n; ++i)
cin>>pigeons[i];
vector<vector<int>> adjList(n);
for(int i = 0; i<n-1; ++i)
{
int a, b;
cin>>a>>b;
adjList[a-1].push_back(b-1);
adjList[b-1].push_back(a-1);
}
init(pigeons, adjList);
cout<<maxPigeonsPath(adjList, pigeons, v)<<endl;
return 0;
}
Compilation message
chase.cpp: In function 'void init(std::vector<long int>&, std::vector<std::vector<int> >&)':
chase.cpp:13:21: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
13 | for(int i = 0; i<pigs.size(); ++i)
| ~^~~~~~~~~~~~
chase.cpp: In function 'long int maxPigeonsPath(std::vector<std::vector<int> >&, std::vector<long int>&, int)':
chase.cpp:60:21: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::vector<int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
60 | for(int i =0 ;i<adj.size(); ++i)
| ~^~~~~~~~~~~
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
0 ms |
212 KB |
Output is correct |
2 |
Incorrect |
0 ms |
212 KB |
Output isn't correct |
3 |
Halted |
0 ms |
0 KB |
- |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
0 ms |
212 KB |
Output is correct |
2 |
Incorrect |
0 ms |
212 KB |
Output isn't correct |
3 |
Halted |
0 ms |
0 KB |
- |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Incorrect |
805 ms |
7404 KB |
Output isn't correct |
2 |
Halted |
0 ms |
0 KB |
- |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
0 ms |
212 KB |
Output is correct |
2 |
Incorrect |
0 ms |
212 KB |
Output isn't correct |
3 |
Halted |
0 ms |
0 KB |
- |