Submission #993838

#	Time	Username	Problem	Language	Result	Execution time	Memory
993838		NValchanov	Wells (CEOI21_wells)	C++17	0 / 100	36 ms	71516 KiB

wells

#include <bits/stdc++.h>

#define endl '\n'

using namespace std;

typedef long long ll;

const int MAXN = 3e6 + 10;
const int MOD = 1e9 + 7;

int n, k;
vector < int > adj[MAXN];
int dist[MAXN];
vector < int > dmt;
int parity[MAXN];

void read()
{
    cin >>  n >> k;
    for(int i = 1; i <= n - 1; i++)
    {
        int u, v;
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
}

void dfs_dist(int u, int par)
{
    dist[u] = dist[par] + 1;

    for(int v : adj[u])
    {
        if(v == par)
            continue;

        dfs_dist(v, u);
    }
}

bool dfs_dmt(int u, int par, int to)
{
    dmt.push_back(u);

    if(u == to)
        return true;
    
    for(int v : adj[u])
    {
        if(v == par)
            continue;
        
        if(dfs_dmt(v, u, to))
            return true;
    }
    dmt.pop_back();
    return false;
}

void find_diameter()
{
    dist[1] = -1;
    dfs_dist(1, 1);

    int end1 = 1;
    for(int i = 1; i <= n; i++)
    {
        if(dist[i] > dist[end1])
            end1 = i;
    }

    dist[end1] = -1;
    dfs_dist(end1, end1);

    int end2 = 1;
    for(int i = 1; i <= n; i++)
    {
        if(dist[i] > dist[end2])
            end2 = i;
    }

    dfs_dmt(end1, end1, end2);
}

void bfs(int start)
{
    for(int i = 1; i <= n; i++)
    {
        parity[i] = -1;
    }

    queue < int > q;
    q.push(start);
    parity[start] = 0;

    while(!q.empty())
    {
        int u = q.front();
        q.pop();

        for(int v : adj[u])
        {
            if(parity[v] == -1)
            {
                parity[v] = (parity[u] + 1) % k;
                q.push(v);
            }
        }
    }
}

vector < int > path;

bool find_path(int u, int par, int to)
{
    path.push_back(u);
    if(u == to)
        return true;

    for(int v : adj[u])
    {
        if(v == par)
            continue;

        if(find_path(v, u, to))
            return true;
    }
    path.pop_back();
    return false;
}

void solve()
{   
    int cnt = 0;
    for(int i = 1; i <= n; i++)
    {
        for(int j = i + 1; j <= n; j++)
        {
            path.clear();
            find_path(i, i, j);

            // cout << "The path from " << i << " to " << j << " : " << endl;

            // for(int u : path)
            // {
            //     cout << u;
            //     if(u != path.back())
            //         cout << " -> ";
            // }
            // cout << endl;

            if(path.size() != k)
                continue;

            unordered_set < int > visited;

            bool lampa = true;

            for(int u : path)
            {
                if(visited.find(parity[u]) != visited.end())
                {
                   lampa = false;
                   break;
                }
                visited.insert(parity[u]);
            }
            cnt += lampa;
        }
    }

    if(cnt)
    {
        cout << "YES" << endl;
        cout << cnt << endl;
    }
    else
    {
        cout << "NO" << endl;
        cout << "0" << endl;    
    }
}

int main()
{
    ios_base :: sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);

    read();
    find_diameter();

    // for(int u : dmt)
    // {
    //     cout << u;
    //     if(u != dmt.back())
    //     {
    //         cout << " -> ";
    //     }
    // }
    // cout << endl;

    bfs(dmt.back());

    // cout << "Parities : " << endl;

    // for(int i = 1; i <= n; i++)
    // {
    //     cout << i << " -> " << parity[i] << endl;
    // }
    // cout << endl;

    solve();

    return 0;
}

Compilation message (stderr)

wells.cpp: In function 'void solve()':
wells.cpp:154:28: warning: comparison of integer expressions of different signedness: 'std::vector<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
  154 |             if(path.size() != k)
      |                ~~~~~~~~~~~~^~~~

• Subtask 1: 0 / 30
• Subtask 2: 0 / 20
• Subtask 3: 0 / 20
• Subtask 4: 0 / 30