[HDU4670]Cube number on a tree
题目
题解
首先，题目中要求我们求的是哪些路径的点积是一个立方数，而每个点的点值都可以被他给定的 \(k\) 个质数组成。
考虑将每个点的点值分解为三进制，第 \(i\) 位表示这个点的点值可以被 \(prime[i]\) 的 \(t\bmod 3\) 次方组成。
至于为什么 \(\bmod 3\) 其实比较显然。
然后，对于每个点，我们定义 map<Ternary,int>MP[i] ，键值 Ternary 即为我们重载的三进制，这个表示以 \(i\) 为根，在它的子树中路径点积为 Ternary 的数量。
但是，由于启发式合并，我们的根有可能被更改，即 MP[i] 有时候不一定对应点 \(i\)，所以我们用一个 rt[i] 表示 \(i\) 真正对应的 MP 到底是哪一个。
启发式合并时，由于我们可能会更改根，涉及到对于 MP 的整个修改，所以我们用一个 tag 作为懒标记，表示整体增加或者减少了多少。
那么，合并之前我们统计答案时，只需要从小的 MP 中找大的 MP 相对应的数值是多少，用组合即数量乘积表示答案。
统计完答案之后，考虑合并 MP，在这之前首先需要修改子树对应的 MP 的 tag，因为合并之后，子树的 MP 应该以当前点为根，而在这之前，子树的 MP 一直以那个子节点为根，所以 tag[rt[v]]+=val[u]。
代码
#include<bits/stdc++.h>
using namespace std;
#define rep(i,__l,__r) for(signed i=(__l),i##_end_=(__r);i<=i##_end_;++i)
#define fep(i,__l,__r) for(signed i=(__l),i##_end_=(__r);i>=i##_end_;--i)
#define writc(a,b) fwrit(a),putchar(b)
#define mp(a,b) make_pair(a,b)
#define ft first
#define sd second
#define LL long long
#define ull unsigned long long
#define uint unsigned int
#define pii pair< int,int >
#define Endl putchar('\n')
// #define int long long
// #define int unsigned
// #define int unsigned long long
#ifdef _GLIBCXX_CSTDIO
#define cg (c=getchar())
template<class T>inline void qread(T& x){
char c;bool f=0;
while(cg<'0'||'9'<c)f|=(c=='-');
for(x=(c^48);'0'<=cg&&c<='9';x=(x<<1)+(x<<3)+(c^48));
if(f)x=-x;
}
template<class T>inline T qread(const T sample){
T x=0;char c;bool f=0;
while(cg<'0'||'9'<c)f|=(c=='-');
for(x=(c^48);'0'<=cg&&c<='9';x=(x<<1)+(x<<3)+(c^48));
return f?-x:x;
}
#undef cg
template<class T>void fwrit(const T x){//just short,int and long long
if(x<0)return (void)(putchar('-'),fwrit(-x));
if(x>9)fwrit(x/10);
putchar(x%10^48);
}
#endif
// template<class T,class... Args>inline void qread(T& x,Args&... args){qread(x),qread(args...);}
template<class T>inline T Max(const T x,const T y){return x>y?x:y;}
template<class T>inline T Min(const T x,const T y){return x<y?x:y;}
template<class T>inline T fab(const T x){return x>0?x:-x;}
inline int gcd(const int a,const int b){return b?gcd(b,a%b):a;}
inline void getInv(int inv[],const int lim,const int MOD){
inv[0]=inv[1]=1;for(int i=2;i<=lim;++i)inv[i]=1ll*inv[MOD%i]*(MOD-MOD/i)%MOD;
}
inline LL mulMod(const LL a,const LL b,const LL mod){//long long multiplie_mod
return ((a*b-(LL)((long double)a/mod*b+1e-8)*mod)%mod+mod)%mod;
}
const int MAXN=50000;
const int MAXK=30;
struct Ternary{
LL num;
Ternary():num(0ll){}
bool operator <(const Ternary rhs)const{
return num<rhs.num;
}
Ternary operator +(const Ternary rhs)const{
LL t=1,f,x1=num,x2=rhs.num;
Ternary ret;
while(x1>0 || x2>0){
f=((x1%3)+(x2%3))%3;
ret.num+=f*t;
x1/=3,x2/=3;
t*=3;
}return ret;
}
Ternary operator -(const Ternary rhs)const{
LL t=1,f,x1=num,x2=rhs.num;
Ternary ret;
while(x1>0 || x2>0){
f=((x1%3)-(x2%3)+3)%3;
ret.num+=f*t;
x1/=3,x2/=3;
t*=3;
}return ret;
}
Ternary operator +=(const Ternary rhs){
return (*this)=(*this)+rhs;
}
Ternary operator -=(const Ternary rhs){
return (*this)=(*this)-rhs;
}
}val[MAXN+5],tag[MAXN+5];//表示单个点的 val 和懒标记
map<Ternary,int>MP[MAXN+5];
LL prime[MAXK+5],ans;
int rt[MAXN+5],n,k;
struct edge{int to,nxt;}e[(MAXN<<1)|2];
int tail[MAXN+5],ecnt;
inline void add_edge(const int u,const int v){
e[++ecnt]=edge{v,tail[u]};tail[u]=ecnt;
e[++ecnt]=edge{u,tail[v]};tail[v]=ecnt;
}
int Merge(int x,int y,const int u){
bool haveswap=false;
if(MP[x].size()<MP[y].size())swap(x,y),haveswap=true;
Ternary tmp;
for(auto it=MP[y].begin();it!=MP[y].end();++it){
tmp=Ternary()-(it->first+tag[y]+tag[x]);
if(MP[x].count(tmp))
ans=ans+1ll*MP[x][tmp]*it->second;
}
if(haveswap)tag[x]+=val[u];
else tag[y]+=val[u];
for(auto it=MP[y].begin();it!=MP[y].end();++it){
tmp=it->first+tag[y]-tag[x];
MP[x][tmp]+=it->second;
}
MP[y].clear();
return x;
}
void dfs(const int u,const int fa){
rt[u]=u;
MP[u][val[u]]=1;
for(int i=tail[u],v;i;i=e[i].nxt)if((v=e[i].to)^fa){
dfs(v,u);
rt[u]=Merge(rt[u],rt[v],u);
}
}
inline void Init(){
while(~scanf("%d %d",&n,&k)){
ecnt=0;ans=0;
rep(i,1,k)prime[i]=qread(1ll);
LL x,cnt,f;
rep(t,1,n){
val[t].num=0,tail[t]=0,tag[t].num=0;
x=qread(1ll),f=1;
rep(i,1,k){cnt=0;
while(x%prime[i]==0){
x/=prime[i];
(++cnt)%=3;
}val[t].num+=cnt*f;
f*=3;
}
if(!val[t].num)++ans;//可以从自己到自己
}
int u,v;
rep(i,1,n-1){
u=qread(1),v=qread(1);
add_edge(u,v);
}
dfs(1,0);
printf("%lld\n",ans);
MP[rt[1]].clear();
}
}
signed main(){
// ios::sync_with_stdio(false);
Init();
return 0;
}
