蚂蚁笔试(二叉树节点的曼哈顿距离)

思路

先造树,再求坐标(由于不确定父子节点的大小关系,可能会先遍历到子节点):哈希存储节点,节点属性包含节点值、坐标、左右孩子指针;按边的前后节点重新排序输入边顺序(便于为父节点确定左右孩子,也可不排序连接父子节点时判断左孩子是否为空);遍历边,构造树;从根节点使用dfs确认节点坐标;求距离。

代码

#include <bits/stdc++.h>
using namespace std;

struct TreeNode{
    int val, x, y;
    TreeNode* left;
    TreeNode* right;
    TreeNode(): val(1), x(0), y(0), left(nullptr), right(nullptr){}
    TreeNode(int value): val(value), left(nullptr), right(nullptr){}
    TreeNode(int value, int x_, int y_): val(value), x(x_), y(y_), left(nullptr), right(nullptr){}
};

bool cmp(vector<int> &a, vector<int>& b){
    return a[0]==b[0]?a[1]<b[1]:a[0]<b[0];
}

int dis(int a, int b, unordered_map<int, TreeNode*>& nodes){
    return abs(nodes[a]->x-nodes[b]->x)+abs(nodes[a]->y-nodes[b]->y);
}

void dfs(TreeNode* node, int x, int y){
    node->x = x;
    node->y = y;
    if(node->left!=nullptr){
        dfs(node->left, x-1, y-1);
    }
    if(node->right!=nullptr){
        dfs(node->right, x+1, y-1);
    }
}

int main() {
    int n, q, i, a, b;
    cin>>n>>q;
    vector<vector<int>> edges(n-1, vector<int>(2, 0));
    unordered_map<int, TreeNode*> nodes;
    TreeNode* tmp;
    for(i=0;i<n-1;i++){
        cin>>edges[i][0]>>edges[i][1];
        if(nodes.find(edges[i][0])==nodes.end()){
            tmp = new TreeNode(edges[i][0]);
            nodes[edges[i][0]] = tmp;
        }
        if(nodes.find(edges[i][1])==nodes.end()){
            tmp = new TreeNode(edges[i][1]);
            nodes[edges[i][1]] = tmp;
        }
    }
    sort(edges.begin(), edges.end(), cmp);
    for(i=0;i<n-1;i++){
        tmp = nodes[edges[i][0]];
        if(tmp->left==nullptr){
            tmp->left = nodes[edges[i][1]];
        }else{
            tmp->right = nodes[edges[i][1]];
        }
    }
    dfs(nodes[1], 0, 0);
    for(;q>0;q--){
        cin>>a>>b;
        cout<<dis(a, b, nodes)<<endl;
    }
    return 0;
}

全部评论

相关推荐

评论
点赞
3
分享

创作者周榜

更多
牛客网
牛客企业服务