E - Subtree K-th Max - Githubissues

https://atcoder.jp/contests/abc239/tasks/abc239_e

問題概要

木にその頂点を示すindexとは別に番号が振られている。 Q回以下の操作を行う。

indexで頂点を指定し、そのindexをルートとした部分木の各頂点に振られている番号のうち、K番目に大きいものを出力せよ。

解き方

各頂点をルートとする部分木の持つ番号を大きい順にソートした配列を作る。ここで肝となるのが、制約が Ki <=20 と小さいこと。そのため20個より大きくなったら切り捨てることでソートのコストがかからない。各頂点をルートとする部分木の持つ番号を大きい順にソートした配列は、帰りがけ順にdfsすることで、子から親にその頂点が持つ番号を伝搬することができる。

type DFSItem struct {
    v int
    parent int
    visitCnt int
}

func solve() {
    in := getInts()
    n, q := in[0], in[1]
    tree := make([][]int, n)
    x := getInts()

    for i := 0; i < n-1; i++ {
        in = getInts()
        a, b := in[0]-1, in[1]-1
        tree[a] = append(tree[a], b)
        tree[b] = append(tree[b], a)
    }

    res := make([][]int, n)
    for i := 0; i < n; i++ {
        res[i] = append(res[i], x[i])
    }

    que := newQueue[*DFSItem]()
    que.push(&DFSItem{v: 0, parent: -1, visitCnt: 1})
    for !que.empty() {
        item := que.pop()
        if item.visitCnt == 1 { // 行きがけ
            que.push(&DFSItem{v: item.v, parent: item.parent, visitCnt: item.visitCnt+1})
            for _, next := range tree[item.v] {
                if next == item.parent {
                    continue
                }
                que.push(&DFSItem{v: next, parent: item.v, visitCnt: 1})
            }
        } else { // 帰りがけ
            for _, next := range tree[item.v] {
                if next == item.parent {
                    continue
                }
                res[item.v] = append(res[item.v], res[next]...)
            }
            res[item.v] = reverse(sortSlice(res[item.v]))
            if len(res[item.v]) > 20 {
                res[item.v] = res[item.v][:20]
            }
        }
    }

    for i := 0; i < q; i++ {
        in = getInts()
        v, d := in[0]-1, in[1]-1
        fmt.Println(res[v][d])
    }
}

https://atcoder.jp/contests/abc239/submissions/53450548

hxrxchang / atcoder

E - Subtree K-th Max #50

問題概要

解き方