5538. 统计子树中城市之间最大距离

2020-10-12  本文已影响0人  来到了没有知识的荒原

5538. 统计子树中城市之间最大距离

状态压缩子树选点
判断子树是否合法
两次bfs求直径

const int N=20;
int dist[N];
queue<int>q;
vector<int> g[N];

void bfs(int root,int s){
    memset(dist,-1,sizeof dist);
    while(q.size())q.pop();
    dist[root]=0;
    q.push(root);
    while(q.size()){
        int now=q.front();q.pop();
        for(auto v:g[now]){
            if (((s >> v) & 1) == 0) continue;
            if(dist[v]==-1){
                dist[v]=dist[now]+1;
                q.push(v);
            }
        }
    }
}

class Solution {
public:
    vector<int> countSubgraphsForEachDiameter(int n, vector<vector<int>>& edges) {
        for(int i=0;i<=n;i++)g[i]=vector<int>();
        vector<int>res(n-1,0);
        
        for(auto v:edges){
            int a=v[0]-1,b=v[1]-1;
            g[a].push_back(b);
            g[b].push_back(a);
        }
        
        int lim=1<<n;
        for(int s=1;s<lim;s++){
            // 遍历所有状态
            int total=0,root=-1;
            for(int i=0;i<n;i++){
                if((s>>i)&1){
                    total++;
                    if(root=-1)root=i;
                }
            }  
            if(total<=1)continue; // 要求是城市对,所以子树中要含2个及以上城市
            
            int mx=root,cnt=0;
            bfs(root,s); // 第一次bfs找到最远点
            
            for(int i=0;i<n;i++){
                if(dist[i]>=0){
                    cnt++;
                    if(dist[i]>dist[mx])mx=i;
                }
            }
            if(cnt!=total)continue; // 判断子树是否成立
            
            bfs(mx,s);  // 从最远点开始第二次bfs找到真实最远点
            
            for(int i=0;i<n;i++){
                if(dist[i]>dist[mx])mx=i;
            }
            
            res[dist[mx]-1]++;
        }
        return res;
    }
};
```![无标题.png](https://img.haomeiwen.com/i11440646/641c6c825509cd6e.png?imageMogr2/auto-orient/strip%7CimageView2/2/w/1240)
上一篇下一篇

猜你喜欢

热点阅读