[Toph] B. Dictionary Tree

 

Problem Link    : B. Dictionary Tree
Category        : Trie, Tree, Dfs
Contest         : Intra AIUB Programming Contest 2024 (Senior)

#include "bits/stdc++.h"
using namespace std;
#define int     long long int
#define endl    '\n'
 
struct Trie {
    int counter;
    Trie *child[26];
    Trie *parent;
 
    Trie() {
        counter = 0;
        parent = nullptr;
        for (int i = 0; i < 26; i++) {
            child[i] = nullptr;
        }
    }
    ~Trie() {
        for (int i = 0; i < 26; i++) {
            if (child[i] and this != child[i]) {
                delete child[i];
            }
        }
    }
};
 
typedef Trie* pnode;
pnode root;
 
signed main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr); cout.tie(nullptr);
    cout.precision(12);
 
    bool FILEIO = true; string FILE = "in.txt";
    if (FILEIO and fopen(FILE.c_str(), "r")) {
        freopen(FILE.c_str(), "r", stdin);
    }
 
    int n, k;
    cin >> n >> k;
    vector<vector<int>> graph(n + 1);
    for (int e = 1; e < n; e++) {
        int u, v;
        cin >> u >> v;
        graph[u].push_back(v);
        graph[v].push_back(u);
    }
    vector<char> letters(n + 1);
    for (int i = 1; i <= n; i++) {
        cin >> letters[i];
    }
    for (int i = 1; i <= n; i++) {
        sort(graph[i].begin(), graph[i].end(), [&](int x, int y)
            {
               return letters[x] < letters[y];
            });
    }
    root = new Trie();
    pnode curRoot = root;
    function<void(int, int)> Dfs = [&](int u, int p) {
        if (curRoot->child[ letters[u] - 'a' ] == nullptr) {
            curRoot->child[ letters[u] - 'a' ] = new Trie();
            curRoot->child[ letters[u] - 'a' ]->parent = curRoot;
            curRoot = curRoot->child[ letters[u] - 'a' ];
            curRoot->counter++;
        } else {
            curRoot->child[ letters[u] - 'a' ]->parent = curRoot;
            curRoot = curRoot->child[ letters[u] - 'a' ];
            curRoot->counter++;
        }
        for (int v : graph[u]) {
            if (v == p) {
                continue;
            }
            Dfs(v, u);
        }
        curRoot = curRoot->parent;
    };
    Dfs(1, 0);
    string path = "";
    function<void(pnode)> Solve = [&](pnode rt) {
        if (rt == nullptr) {
            return;
        }
        if (k <= 0) {
            cout << path << endl;
            exit(0);
        }
        for (int i = 0; i < 26; i++) {
            if (rt->child[i] == nullptr) {
                continue;
            }
            path += static_cast<char>(i + 'a');
            k -= rt->child[i]->counter;
            Solve(rt->child[i]);
            path.pop_back();
        }
    };
    Solve(root);
}