Submission #514838

#	Time	Username	Problem	Language	Result	Execution time	Memory
514838		blue	Wells (CEOI21_wells)	C++17	0 / 100	156 ms	311560 KiB

wells

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

//one-indexed!!!

const int maxN = 1'500'000;
const int logN = 22;

int N, K;

vector<int> a(maxN);
vector<int> b(maxN);

vector<int> edge[1+maxN];
vector<int> children[1+maxN];





const int MINVALUE = -1e9; //?????????????
struct segtree
{
    int l;
    int r;

    int lp = 0;
    int mx = 0;

    segtree* left = NULL;
    segtree* right = NULL;

    segtree()
    {
        ;
    }

    segtree(int L, int R)
    {
        l = L;
        r = R;
        if(l == r) return;
        int m = (l+r)/2;
        left = new segtree(l, m);
        right = new segtree(m+1, r);
    }

    void add(int L, int R, int V)
    {
        if(R < l || r < L || R < L) return;
        else if(L <= l && r <= R)
        {
            lp += V;
            mx += V;
        }
        else
        {
            left->add(L, R, V);
            right->add(L, R, V);
            mx = lp + max(left->mx, right->mx);
        }
    }

    int rangemax(int L, int R)
    {
        if(R < l || r < L || R < L) return MINVALUE;
        else if(L <= l && r <= R) return mx;
        else return lp + max(left->rangemax(L, R), right->rangemax(L, R));
    }
};









//Precompute
int curr_ind = 0;
vector<int> order(1);
vector<int> order_index(1+maxN, -1); //direct
vector<int> depth(1+maxN, -1); //give to child

//Postcompute
vector<int> subtree(1+maxN, 1); //initialized to 1!!!!!
int anc[1+maxN][1+logN];

void dfs1(int u)
{
    curr_ind++;

    order.push_back(u);
    order_index[u] = curr_ind;

    for(int v: edge[u])
    {
        if(order_index[v] != -1) continue;

        depth[v] = depth[u] + 1;
        children[u].push_back(v);
        anc[v][0] = u;

        dfs1(v);

        subtree[u] += subtree[v];
    }
}




vector<int> good(1+maxN, 0); //Is the node good? (Is there any path of length K which contains this node?)
segtree S;

void dfs2(int u) //u's processing is done before dfs2(u) is called
{
    bool is_good = 0;
    if(S.rangemax(1, N) >= K-1)
        is_good = 1;
    else
    {
        vector<int> subtree_depths;
        for(int v: children[u]) //children
        {
            subtree_depths.push_back(S.rangemax(order_index[v], order_index[v] + subtree[v] - 1));
        }

        //remainder of the tree
        if(u != 1)
        {
            int L = S.rangemax(1, order_index[u] - 1); //returns MINVALUE if the interval is empty.

            int R = S.rangemax(order_index[u] + subtree[u], N);

            subtree_depths.push_back(max(L, R));
        }
        sort(subtree_depths.begin(), subtree_depths.end());

        int DS = (int)subtree_depths.size();

        if(DS >= 2)
        {
            if(subtree_depths[DS-1] + subtree_depths[DS-2] + 1 >= K)
                is_good = 1;
        }
    }

    good[u] = is_good;






    for(int v: children[u])
    {
        S.add(1, N, +1);
        S.add(order_index[v], order_index[v] + subtree[v] - 1, -2);

        dfs2(v);

        S.add(1, N, -1);
        S.add(order_index[v], order_index[v] + subtree[v] - 1, +2);
    }
}







const long long mod = 1'000'000'007LL;

int curr_comp = 0;
vector<int> comp(1+maxN, 0);
vector<int> comp_list[1+maxN];

vector<int> newEdge[1+maxN];

void dfs3(int u)
{
    comp[u] = curr_comp;
    comp_list[ comp[u] ].push_back(u);
    for(int v: edge[u])
    {
        if(!good[v]) continue;
        if(comp[v]) continue;
        dfs3(v);
    }
}


int lca(int u, int v)
{
    if(depth[u] > depth[v]) swap(u, v);
    for(int e = logN; e >= 0; e--)
    {
        if(depth[v] - (1 << e) < depth[u]) continue;
        v = anc[v][e];
    }
    if(u == v) return u;

    for(int e = logN; e >= 0; e--)
    {
        if(anc[u][e] != anc[v][e])
        {
            u = anc[u][e];
            v = anc[v][e];
        }
    }
    u = anc[u][0];
    v = anc[v][0];
    return u;
}

int dist(int u, int v)
{
    return depth[u] + depth[v] - 2*depth[lca(u, v)];
}










const int nopath = -1;

vector<int> visit(1+maxN, 0); //be careful while using this!
vector<int> pathpos(1+maxN, nopath);

vector<int> path;
void pathfinder(int u)
{
    // cerr << "pathfinder " << u << '\n';
    // cerr << K << '\n';
    visit[u] = 1;
    path.push_back(u);
    if((int)path.size() == K) return;

    for(int v: newEdge[u])
    {
        if(visit[v]) continue;
        pathfinder(v);
        if((int)path.size() == K) return;
    }
    path.pop_back();
}

vector<int> new_depth(1+maxN, -1);
vector< vector<int> > desc_list[1+maxN];
int new_anc[1+maxN][1+logN];

int p;

int ordercounter = 0;
vector<int> orderpos(1+maxN, -1);

void add_descendant(int P, int u, int d)
{
    while((int)desc_list[P].size() - 1 < d)
        desc_list[P].push_back(vector<int>(0));

    desc_list[P][d].push_back(u);
}

void new_dfs(int P, int u)
{
    for(int v: newEdge[u])
    {
        if(new_depth[v] != -1) continue; //v is not already visited
        if(pathpos[v] != nopath) continue; //v is not on the path.

        new_depth[v] = 1 + new_depth[u];
        add_descendant(P, v, new_depth[v]);

        ordercounter++;
        orderpos[v] = ordercounter;

        new_anc[v][0] = u;
        new_dfs(P, v);
    }
}

vector<int> sourcetree[1+maxN];
vector<int> listofdesc[1+maxN];

int currP;
vector<int> path1;

int new_lca(int u, int v)
{
    if(new_anc[u][logN] != new_anc[v][logN])
    {
        u = new_anc[u][logN];
        v = new_anc[v][logN];
        if(pathpos[u] == currP || pathpos[v] == currP) return u;


        if(pathpos[u] > pathpos[v]) swap(u, v);

        if(max(pathpos[u], pathpos[v]) < currP) return v;
        else if(min(pathpos[u], pathpos[v]) > currP) return u;
        else return path1[currP];
    }
    else
    {
        if(new_depth[u] > new_depth[v]) swap(u, v);
        for(int e = logN; e >= 0; e--)
        {
            if(new_depth[v] - (1 << e) < new_depth[u]) continue;
            v = new_anc[v][e];
        }
        if(u == v) return u;

        for(int e = logN; e >= 0; e--)
        {
            if(new_anc[u][e] != new_anc[v][e])
            {
                u = new_anc[u][e];
                v = new_anc[v][e];
            }
        }
        u = new_anc[u][0];
        v = new_anc[v][0];
        return u;
    }
}



long long solve(int c)
{
    path.clear();

    int src = 0;
    for(int u: comp_list[c])
        if((int)newEdge[u].size() == 1)
        {
            src = u;
            break;
        }

    long long res = 0;
    // cerr << "src = " << src << '\n';

    pathfinder(src);
    for(int i = 0; i < K; i++)
        pathpos[ path[i] ] = i;
    path1 = path;

    for(int p: path)
    {
        new_depth[p] = 0;

        desc_list[p].push_back(vector<int>(0));
        desc_list[p][0].push_back(p);

        new_anc[p][0] = p;

        ordercounter = 0;
        new_dfs(p, p);
    }

    for(int e = 1; e <= logN; e++)
    {
        for(int u: comp_list[c])
        {
            new_anc[u][e] = new_anc[ new_anc[u][e-1] ][e-1];
        }
    }

    for(int f: comp_list[c])
    {
        if(pathpos[new_anc[f][logN]] + K - (new_depth[f] % K) < K)
            sourcetree[f].push_back(path[ pathpos[new_anc[f][logN]] + K - (new_depth[f] % K)]);

        if(0 <= pathpos[new_anc[f][logN]] - K + (new_depth[f] % K))
            sourcetree[f].push_back(path[ pathpos[new_anc[f][logN]] - K + (new_depth[f] % K)]);

        if(new_depth[f] == 0)
            sourcetree[f].push_back(f);
    }

    for(int f: comp_list[c])
    {
        for(int st: sourcetree[f])
        {
            listofdesc[st].push_back(f);
        }
    }

    for(int p1: path)
    {
        p = p1;
        sort(listofdesc[p].begin(), listofdesc[p].end(), [] (int x, int y)
        {
            if(abs(pathpos[new_anc[x][logN]] - pathpos[p]) != abs(pathpos[new_anc[y][logN]] - pathpos[p]) )
            {
                return abs(pathpos[new_anc[x][logN]] - pathpos[p]) < abs(pathpos[new_anc[y][logN]] - pathpos[p]);
            }

            if(pathpos[new_anc[x][logN]] != pathpos[new_anc[y][logN]])
                return pathpos[new_anc[x][logN]] < pathpos[new_anc[y][logN]];

            return orderpos[x] < orderpos[y];
        });

        currP = pathpos[p1];


        bool works = 1;
        for(int i = 0; i+1 < listofdesc[p].size(); i++)
        {
            int A = listofdesc[p][i];
            int B = listofdesc[p][i+1];

            int C = new_lca(A, B);
            if(dist(C, src) % (K-1) == 0)
            {
                ;
            }
            else if((K-1)%2 == 0 && dist(C, src) % ((K-1)/2) == 0)
            {
                ;
            }
            else
            {
                works = 0;
            }
        }

        if(works)
        {
            // cerr << "p = " << p1 << '\n';
            res++;
        }
    }

    return res;
}






int main()
{



//PART 0: Input
    cin >> N >> K;
    for(int e = 1; e <= N-1; e++)
    {
        cin >> a[e] >> b[e];
        edge[ a[e] ].push_back(b[e]);
        edge[ b[e] ].push_back(a[e]);
    }






//PART 1: Dividing the tree into components



//1A. Make the segment tree
    S = segtree(1, N);



//1B. Preliminary DFS (dfs1)
    anc[1][0] = 1;
    depth[1] = 0;
    dfs1(1);
    for(int e = 1; e <= logN; e++)
    {
        for(int i = 1; i <= N; i++)
        {
            anc[i][e] = anc[ anc[i][e-1] ][e-1];
        }
    }


//1C. Compute the list of good nodes
    for(int i = 1; i <= N; i++)
        S.add(order_index[i], order_index[i], +depth[i]);
    dfs2(1);


//1D. Divide up the components.
    for(int i = 1; i <= N; i++)
    {
        if(comp[i]) continue;
        if(!good[i]) continue;
        curr_comp++;
        dfs3(i);
    }
    for(int e = 1; e <= N-1; e++)
    {
        if(comp[ a[e] ] == comp[ b[e] ])
        {
            newEdge[ a[e] ].push_back( b[e] );
            newEdge[ b[e] ].push_back( a[e] );
        }
    }











//PART 2: Compute the answer for each component.

    long long res = 1;

    for(int c = 1; c <= curr_comp; c++)
        res = (res * solve(c)) % mod;

    // for(int i = 1; i <= N; i++) cerr << good[i] << ' ';
    // cerr << '\n';

    for(int i = 1; i <= N; i++)
        if(!good[i])
            res = (res * 2) % mod;


    if(res == 0)
        cout << "NO\n";
    else
        cout << "YES\n";

    cout << res << '\n';
}

Compilation message (stderr)

wells.cpp: In function 'long long int solve(int)':
wells.cpp:421:28: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  421 |         for(int i = 0; i+1 < listofdesc[p].size(); i++)
      |                        ~~~~^~~~~~~~~~~~~~~~~~~~~~

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