Capital Movement | Competitive Programming

Question

PROBLEM

Suppose there are 'n' cities in a country out of which, one is the capital. All of the cities are connected through 'n-1' roads that it's possible to travel between any two cities using those roads. Population of cities P_i is provided in the question.

Now if a city is infected with virus, the cities connected with that city also get infected. In the case of an infection a new capital is selected. The non-infected city having the largest population becomes the new capital.

INPUT

The first line of the input contains an integer T denoting the number of test cases. The description of T test cases follows.

The first line of each test case contains one integer N which denotes number of cities.

Next line contains N space-separated integers P1, P2, ..., PN denoting the population of each city.

Next N-1 lines contain two space-separated integers each V and U denoting that there's a road between cities V and U.

Output

For each test case, output a single line containing N integers A1, A2, ..., AN separated by a space. Here Ai denotes the number of the city picked to be new capital in case infection starts spreading from the city i. In case infection affects all the cities output 0.

Example

Input:

1

6

5 10 15 20 25 30

1 3

2 3

3 4

4 5

4 6

Output:

6 6 6 2 6 5

MY C++ SOLUTION

#include <iostream>
#include <vector>
#include <list>
#include <stdio.h>
#include <algorithm>
using namespace std;

int main() {
    int t;
    scanf("%d", &t); //No. of test cases
    while(t--) {
        int n;
        scanf("%d", &n);
        vector<int> p(n);   //vector for storing population of each city
        for(int i = 0; i < n; i++)
            scanf("%d", &p[i]);

        vector<list<int> > graph(n); //vector of lists for storing connection between cities


        for(int i = 0; i < n-1; i++) { 
            int a, b;
            scanf("%d %d", &a, &b);
            graph[--a].push_back(--b);
            graph[b].push_back(a);
        }

        for(int i = 0; i < n; i++) {
            vector<int> temp = p;             //Temporary vector 

        /*All the infected cities are assigned a population
          of -1 so that they don't compete to become the largest capital*/

            temp[i] = -1;
            if(graph[i].size() == n-1) {//Printing "0" if all cities have connection with the infected city.
                cout<<"0 ";
                continue;
            }
            int s = graph[i].size();
            for(int j = 0; j < s; j++) {
                temp[graph[i].front()] = -1;
                graph[i].pop_front();
            }

            /*Finding max. population*/
            int maxindex = std::distance(temp.begin(), std::max_element(temp.begin(), temp.end()));

            printf("%d ", maxindex+1);
        }
        printf("\n");
    }
    return 0;
} 

The function max_element is making the time complexity quadratic. My solution is getting time limit exceeded for large inputs. Hence, my program needs improvement in time taken on execution. Please help. Thanks for your time.

Answer 1

Solved! The trick was to sort the population of cities beforehand.

#include <iostream>
#include <vector>
#include <list>
#include <stdio.h>
#include <algorithm>
using namespace std;
bool myfnc(pair<int, int> i, pair<int, int> j) {
    return i.first > j.first;
}
int main() {
    int t;
    scanf("%d", &t);
    while(t--) {
        int n, d;
        scanf("%d", &n);
        vector<int> p(n);
        int m = 0, mi = 0;
        for(int i = 0; i < n; i++) {
            scanf("%d", &d);
            p[i] = d;
        }
        vector<pair<int, int> > p1(n);
        for(int i = 0; i < n; i++) {
            p1[i].first = p[i];
            p1[i].second = i;
        }
        sort(p1.begin(), p1.end(), myfnc);

        vector<vector<bool> > graph(n, vector<bool> (n));
        for(int i = 0; i < n-1; i++) {
            int a, b;
            scanf("%d %d", &a, &b);
            graph[--a][--b] = 1;
            graph[b][a] = 1;
            graph[i][i] = 1;
        }
        graph[n-1][n-1] = 1;
        for(int i = 0; i < n; i++) {
            int j;
            for(j = 0; j < n; j++) {
                if(graph[i][p1[j].second] == 0) {
                    printf("%d ", p1[j].second+1);
                    break;
                }
            }
            if(j == n)
                printf("0 ");
        }
        printf("\n");
    }
    return 0;
}

Answer 2

Because there is no information on what "bigger inputs" means I'm guessing it refers to having more cities. In that case, you are making a huge (nxn) graph array which will cause more memory to be hogged while at the same time being (mostly) empty as there are only n-1 roads. Consider connecting the cities with std::list and then look only for cities not in the same list as the infected city.