牛客练习赛104 F

Update on 2022.10.22：

出题人已经修改数据并重测了。

Update on 2022.10.21：

实践检验表明数据确实锅了，对 $k=1k=1$ 的数据请输出 $00$！

完蛋了交了 $101101$ 发要被封号了。

不管了，先胡一下这个过不去的做法，供后人观赏。

倒开人表示很亏啊，被这题浪费了这么久……

如果正解是什么高明稳定正确做法当我没说。

考虑 hash。

我们把每个数往 $[0,2^{64})[0,2^\{64\})$ 上做随机映射，取和自然溢出。

然后，淀粉质维护信息，直接用 map 更新答案。

这个复杂度是 $O(n{}^{2}n)O(n\backslash log^2n)$ 的，TLE 过不去。

换成高速的 gp_hash_table hash 表，然后双哈希，会 TLE，单哈希会 WA。

把 vector 存图改成前向星，双哈希会随机 WA/TLE，三哈希会 TLE。

听说有位常数小的老哥写了四哈希，还是 WA 了。

。。。

但是准确性理论上还是可以保证的。

值得注意的是，如果写的是异或哈希，你必须保证两条贡献的链分别得不重复、不赘余进行枚举，因此其实不如直接写相加的哈希。

下面是过不去的代码。我把随机种子改了交了几十发过不去。

（听 SH 说数据有问题，咱也不知道，咱也不敢问。）

// 那就是希望。
// 即便需要取模，也是光明。

// 简单 hash。
// O(n\log^2n)

#ifndef MYEE
#define _CRT_SECURE_NO_WARNINGS
#pragma GCC optimize(2)
#pragma GCC optimize(3)
#pragma GCC optimize("Ofast")
#pragma GCC optimize("inline")
#pragma GCC optimize("-fgcse")
#pragma GCC optimize("-fgcse-lm")
#pragma GCC optimize("-fipa-sra")
#pragma GCC optimize("-ftree-pre")
#pragma GCC optimize("-ftree-vrp")
#pragma GCC optimize("-fpeephole2")
#pragma GCC optimize("-ffast-math")
#pragma GCC optimize("-fsched-spec")
#pragma GCC optimize("unroll-loops")
#pragma GCC optimize("-falign-jumps")
#pragma GCC optimize("-falign-loops")
#pragma GCC optimize("-falign-labels")
#pragma GCC optimize("-fdevirtualize")
#pragma GCC optimize("-fcaller-saves")
#pragma GCC optimize("-fcrossjumping")
#pragma GCC optimize("-fthread-jumps")
#pragma GCC optimize("-funroll-loops")
#pragma GCC optimize("-fwhole-program")
#pragma GCC optimize("-freorder-blocks")
#pragma GCC optimize("-fschedule-insns")
#pragma GCC optimize("inline-functions")
#pragma GCC optimize("-ftree-tail-merge")
#pragma GCC optimize("-fschedule-insns2")
#pragma GCC optimize("-fstrict-aliasing")
#pragma GCC optimize("-fstrict-overflow")
#pragma GCC optimize("-falign-functions")
#pragma GCC optimize("-fcse-skip-blocks")
#pragma GCC optimize("-fcse-follow-jumps")
#pragma GCC optimize("-fsched-interblock")
#pragma GCC optimize("-fpartial-inlining")
#pragma GCC optimize("no-stack-protector")
#pragma GCC optimize("-freorder-functions")
#pragma GCC optimize("-findirect-inlining")
#pragma GCC optimize("-fhoist-adjacent-loads")
#pragma GCC optimize("-frerun-cse-after-loop")
#pragma GCC optimize("inline-small-functions")
#pragma GCC optimize("-finline-small-functions")
#pragma GCC optimize("-ftree-switch-conversion")
#pragma GCC optimize("-foptimize-sibling-calls")
#pragma GCC optimize("-fexpensive-optimizations")
#pragma GCC optimize("-funsafe-loop-optimizations")
#pragma GCC optimize("inline-functions-called-once")
#pragma GCC optimize("-fdelete-null-pointer-checks")
#endif

#include <algorithm>
#include <map>
#include <random>
#include <stdio.h>
#include <vector>
#include <ext/pb_ds/hash_policy.hpp>
#include <ext/pb_ds/assoc_container.hpp>
using namespace std;
using namespace __gnu_pbds;
typedef long long llt;
typedef unsigned uint;typedef unsigned long long ullt;
typedef bool bol;typedef char chr;typedef void voi;
typedef double dbl;
template<typename T>bol _max(T&a,T b){return(a<b)?a=b,true:false;}
template<typename T>bol _min(T&a,T b){return(b<a)?a=b,true:false;}
template<typename T>T lowbit(T n){return n&-n;}
template<typename T>T gcd(T a,T b){return b?gcd(b,a%b):a;}
template<typename T>T lcm(T a,T b){return(a!=0||b!=0)?a/gcd(a,b)*b:(T)0;}
template<typename T>T exgcd(T a,T b,T&x,T&y){if(b!=0){T ans=exgcd(b,a%b,y,x);y-=a/b*x;return ans;}else return y=0,x=1,a;}
template<typename T>T power(T base,T index,T mod)
{
    T ans=1%mod;
    while(index)
    {
        if(index&1)ans=ans*base%mod;
        base=base*base%mod,index>>=1;
    }
    return ans;
}
struct my_hash {
  static ullt splitmix64(ullt x) {
    x += 0x9e3779b97f4a3c15;
    x = (x ^ (x >> 30)) * 0xbf58473d1ce4e5b9;
    x = (x ^ (x >> 27)) * 0x94d049bb133111eb;
    return x ^ (x >> 31);
  }
  size_t operator()(pair<ullt, ullt> x) const {
    static const ullt FIXED_RANDOM =8327984759743698;
    return splitmix64(x.first + FIXED_RANDOM) ^
           (splitmix64(x.second + FIXED_RANDOM) >> 1);
  }
};
std::mt19937_64 rng(233);
uint W[300005];
ullt ans,To[2][300005],All[2];
gp_hash_table<std::pair<ullt,ullt>,ullt,my_hash>M;
uint Last[600005],T[600005],End[300005],tp;
bol Gone[300005],Init[300005];
uint getsiz(uint p,uint f){
    uint siz=1;
    for(uint e=End[p],s;~e;e=Last[e])if(!Gone[s=T[e]]&&s!=f)siz+=getsiz(s,p);
    return siz;
}
uint getwp(uint p,uint f,uint siz,uint&w,uint&wp){
    uint ans=1,t=0,user;
    for(uint e=End[p],s;~e;e=Last[e])if(!Gone[s=T[e]]&&s!=f)ans+=user=getwp(s,p,siz,w,wp),_max(t,user);
    _max(t,siz-ans);if(_min(w,t))wp=p;
    return ans;
}
uint k;
voi dfs1(uint p,uint f,ullt A,ullt B){
    if(W[p]>k||Init[W[p]])return;
    Init[W[p]]=1;ans+=M[{A+=To[0][W[p]],B+=To[1][W[p]]}];
    for(uint e=End[p],s;~e;e=Last[e])if(!Gone[s=T[e]]&&s!=f)dfs1(s,p,A,B);
    Init[W[p]]=0;
}
voi dfs2(uint p,uint f,ullt A,ullt B){
    if(W[p]>k||Init[W[p]])return;
    Init[W[p]]=1;M[{A-=To[0][W[p]],B-=To[1][W[p]]}]++;
    for(uint e=End[p],s;~e;e=Last[e])if(!Gone[s=T[e]]&&s!=f)dfs2(s,p,A,B);
    Init[W[p]]=0;
}
voi dfs(uint p){
    {uint n=getsiz(p,-1);getwp(p,-1,n,n,p);}
    Gone[p]=true;
    if(W[p]<=k){
        M.clear(),M[{All[0],All[1]}]=1;
        Init[W[p]]=1;
        for(uint e=End[p],s;~e;e=Last[e])if(!Gone[s=T[e]])
            dfs1(s,p,To[0][W[p]],To[1][W[p]]),dfs2(s,p,All[0],All[1]);
        Init[W[p]]=0;
    }
    for(uint e=End[p],s;~e;e=Last[e])if(!Gone[s=T[e]])dfs(s);
}
voi insert(uint u,uint v){Last[tp]=End[u],T[tp]=v,End[u]=tp++;}
int main()
{
#ifdef MYEE
    freopen("QAQ.in","r",stdin);
#endif
    uint n;scanf("%u%u",&n,&k);
    for(uint i=0;i<n;i++)scanf("%u",W+i),End[i]=-1;
    if(k==1){
        uint ans=0;for(uint i=0;i<n;i++)ans+=W[i]==1;
        printf("%u\n",ans);return 0;
    }
    for(uint i=0;i<=n;i++)To[0][i]=rng(),To[1][i]=rng();
    for(uint i=1;i<=k;i++)All[0]+=To[0][i],All[1]+=To[1][i];
    for(uint i=1,u,v;i<n;i++)scanf("%u%u",&u,&v),insert(--u,--v),insert(v,u);
    dfs(0),printf("%llu\n",ans);
    return 0;
}

// 那就是希望。
// 即便需要取模，也是光明。