Scheduling Problem on Forest
(src/tree/ForestScheduling.hpp)

• View this file on GitHub
• Last update: 2023-06-17 00:18:57+09:00
• Include: #include "src/tree/ForestScheduling.hpp"

概要

$n$ 頂点の根付き森が与えられ，頂点を先祖が子孫の頂点の右側に置かれないように並べ替えるときに頂点に書かれた値の列のスコアを最大（最小）化する問題を $\mathrm{O}(n \log n)$ 時間で解く． この際，値の列の優先関係について，反射律と推移律，完全律が成立する弱順序となっている必要がある． S に単位元と < の演算子，マージ関数 op が定義されている必要がある．

出題例

• AtCoder Grand Contest 023 F - 01 on Tree
• Toyota Programming Contest 2023 Spring Final G - Git Gud
• The 1st Universal Cup. Stage 19: America D. Fail Fast

Code

#pragma once
#include <cassert>
#include <queue>
#include <vector>

template <class S, S (*op)(S, S)> struct ForestScheduling {
    S tot;
    std::vector<int> ord;

    ForestScheduling(const std::vector<int>& par, std::vector<S> x) : tot() {
        int n = par.size();
        assert(int(x.size()) == n);
        dsu.assign(n + 1, -1);
        std::priority_queue<T> pq;
        std::vector<int> tails(n + 1), nxt(n + 1, -1);
        for (int i = 0; i < n; i++) {
            tails[i] = i;
            pq.push({x[i], i, i});
        }
        tails[n] = n;
        while (not pq.empty()) {
            auto [tmp, v, tail] = pq.top();
            pq.pop();
            if (tails[v] == tail) {
                int u = leader(par[v] >= 0 ? par[v] : n);
                merge(u, v);
                nxt[tails[u]] = v;
                tails[u] = tail;
                if (u == n)
                    tot = op(tot, x[v]);
                else {
                    x[u] = op(x[u], x[v]);
                    pq.push({x[u], u, tail});
                }
            }
        }
        for (int i = 0, cur = n; i < n; i++) {
            cur = nxt[cur];
            ord.emplace_back(cur);
        }
    }

    int operator[](int i) const { return ord[i]; }

  private:
    struct T {
        S s;
        int head, tail;
        bool operator<(const T& rhs) const { return rhs.s < s; }
    };
    std::vector<int> dsu;

    int leader(int u) { return dsu[u] < 0 ? u : (dsu[u] = leader(dsu[u])); }

    bool merge(int u, int v) {
        u = leader(u), v = leader(v);
        if (u == v) return false;
        dsu[u] += dsu[v];
        dsu[v] = u;
        return true;
    }
};

#line 2 "src/tree/ForestScheduling.hpp"
#include <cassert>
#include <queue>
#include <vector>

template <class S, S (*op)(S, S)> struct ForestScheduling {
    S tot;
    std::vector<int> ord;

    ForestScheduling(const std::vector<int>& par, std::vector<S> x) : tot() {
        int n = par.size();
        assert(int(x.size()) == n);
        dsu.assign(n + 1, -1);
        std::priority_queue<T> pq;
        std::vector<int> tails(n + 1), nxt(n + 1, -1);
        for (int i = 0; i < n; i++) {
            tails[i] = i;
            pq.push({x[i], i, i});
        }
        tails[n] = n;
        while (not pq.empty()) {
            auto [tmp, v, tail] = pq.top();
            pq.pop();
            if (tails[v] == tail) {
                int u = leader(par[v] >= 0 ? par[v] : n);
                merge(u, v);
                nxt[tails[u]] = v;
                tails[u] = tail;
                if (u == n)
                    tot = op(tot, x[v]);
                else {
                    x[u] = op(x[u], x[v]);
                    pq.push({x[u], u, tail});
                }
            }
        }
        for (int i = 0, cur = n; i < n; i++) {
            cur = nxt[cur];
            ord.emplace_back(cur);
        }
    }

    int operator[](int i) const { return ord[i]; }

  private:
    struct T {
        S s;
        int head, tail;
        bool operator<(const T& rhs) const { return rhs.s < s; }
    };
    std::vector<int> dsu;

    int leader(int u) { return dsu[u] < 0 ? u : (dsu[u] = leader(dsu[u])); }

    bool merge(int u, int v) {
        u = leader(u), v = leader(v);
        if (u == v) return false;
        dsu[u] += dsu[v];
        dsu[v] = u;
        return true;
    }
};

Back to top page