#724573 (C++17(gcc12)) No.1784 Not a star yet...

提出ソース
結果

問題	No.1784 Not a star yet...
コンテスト
ユーザー	Nachia
提出日時	2021-12-17 01:53:12
言語	C++17(gcc12) (gcc 12.3.0 + boost 1.87.0)
結果	AC
実行時間	64 ms / 2,000 ms
コード長	10,101 bytes
コンパイル時間	6,639 ms
コンパイル使用メモリ	138,712 KB
最終ジャッジ日時	2025-01-27 01:43:25
ジャッジサーバーID （参考情報）	judge4 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン	結果
sample	AC * 3
other	AC * 61
権限があれば一括ダウンロードができます
ソースコード

raw source code

#include <vector>
#include <cassert>
#include <utility>
#include <atcoder/modint>


template<class Elem, atcoder::internal::is_static_modint_t<Elem>* = nullptr>
struct modulo_matrix{
private:
    using m32 = Elem;
public:
    int h;
    int w;
    std::vector<m32> elems;
    
    modulo_matrix(int new_h, int new_w){ h = new_h; w = new_w; elems.assign(h * w, 0); }
    modulo_matrix(const modulo_matrix&) = default;
    typename std::vector<m32>::iterator operator[](int y){ return elems.begin() + (y * w); }
    typename std::vector<m32>::const_iterator operator[](int y) const { return elems.begin() + (y * w); }
    static modulo_matrix identity(int idx){ auto res = modulo_matrix(idx, idx); for(int i = 0; i < idx; i++) res[i][i] = 1; return res; }
    modulo_matrix operator*(const modulo_matrix& r) const {
        assert(w == r.h);
        auto res = modulo_matrix(h, r.w);
        for (int i=0; i<h; i++) for (int j=0; j<w; j++) for (int k=0; k<r.w; k++) res[i][k] += (*this)[i][j] * r[j][k];
        return res;
    }
    m32 det() const {
        assert(h == w);
        modulo_matrix g = *this;
        m32 ans = 1;
        for (int i=0; i<h; i++) {
            int tg = -1;
            for (int j=i; j<h; j++) { if (g[j][i] != 0) tg = j; }
            if (tg == -1) return 0;
            if (tg != i) ans = -ans;
            for (int j=0; j<h; j++) std::swap(g[i][j], g[tg][j]); tg = i;
            ans *= g[i][i];
            m32 const_coeff = g[i][i].inv();
            for (int j=0; j<h; j++) g[i][j] *= const_coeff;
            for (int j=i+1; j<h; j++) for(int k=h-1; k>=i; k--) g[j][k] -= g[j][i] * g[i][k];
        }
        return ans;
    }
    int rank() const {
        modulo_matrix g = *this;
        int y = 0;
        for (int d=0; d<w; d++) {
            if(y == h) break;
            int tg = -1;
            for (int i=y; i<h; i++) { if (g[i][d] != 0){ tg = i; break; } }
            if (tg == -1) continue;
            for (int j=d; j<w; j++) std::swap(g[y][j], g[tg][j]); tg = y;
            m32 const_coeff = g[y][d].inv();
            for (int j=d; j<w; j++) g[y][j] *= const_coeff;
            for (int i=y+1; i<h; i++) for(int j=w-1; j>=d; j--) g[i][j] -= g[i][d] * g[y][j];
            y++;
        }
        return y;
    }
    std::pair<std::vector<m32>, modulo_matrix> linear_equation() const {
        modulo_matrix g = *this;
        int y = 0;
        std::vector<std::pair<int,int>> det_var;
        std::vector<int> rank_var;
        for (int d=0; d<w-1; d++) {
            int tg = -1;
            for (int i=y; i<h; i++) { if (g[i][d] != 0){ tg = i; break; } }
            if (tg == -1){ rank_var.push_back(d); continue; }
            for (int j=d; j<w; j++) std::swap(g[y][j], g[tg][j]); tg = y;
            m32 const_coeff = g[y][d].inv();
            for (int j=d; j<w; j++) g[y][j] *= const_coeff;
            for (int i=0; i<h; i++) if (i != y) for(int j=w-1; j>=d; j--) g[i][j] -= g[i][d] * g[y][j];
            det_var.push_back(std::make_pair(d,y));
            y++;
        }
        for (int i=y; i<h; i++) if (g[i][w-1].val() != 0) return std::make_pair(std::vector<m32>(), modulo_matrix(0,0));
        std::vector<m32> solution(w-1, 0);
        for (auto [x,i] : det_var) { solution[x] = g[i][w-1]; }
        modulo_matrix solution_rank(rank_var.size(), w-1);
        for (int d=0; d<rank_var.size(); d++) {
            int varid = rank_var[d];
            solution_rank[d][varid] = -1;
            for (auto [x,i] : det_var) { solution_rank[d][x] = g[i][varid]; }
        }
        return std::make_pair(std::move(solution), std::move(solution_rank));
    }
};

#include <iostream>
#include <algorithm>
using namespace std;
using i32 = int32_t;
using u32 = uint32_t;
using i64 = int64_t;
using u64 = uint64_t;
#define rep(i,n) for(int i=0; i<(n); i++)

using m32 = atcoder::static_modint<998244353>;
using Matrix = modulo_matrix<m32>;

struct EquationTerm{
    int idx_1;
    int idx_2;
    m32 coeff;
};

// ... = 0
struct Equation{
    vector<EquationTerm> terms;
    m32 constant_term;
    m32 get_coeff(int idx_1, int idx_2){
        for(auto& t : terms){
            if(t.idx_1 == idx_1 && t.idx_2 == idx_2) return t.coeff;
        }
        return 0;
    }
    void add(EquationTerm term){
        for(auto& t : terms){
            if(t.idx_1 == term.idx_1 && t.idx_2 == term.idx_2){
                t.coeff += term.coeff;
                if(t.coeff.val() == 0){
                    swap(t, terms.back());
                    terms.pop_back();
                    return;
                }
                return;
            }
        }
        terms.push_back(term);
    }
    void output(){
        rep(i,terms.size()){
            cout << terms[i].coeff.val() << " f( " << terms[i].idx_1 << ", " << terms[i].idx_2 << " ) + ";
        }
        cout << constant_term.val() << " = 0" << endl;
    }
};


int N;
int num_edge1;
int num_edge2;
vector<pair<int,int>> initial_state;
m32 invN;
m32 invWeight;
m32 invNumOpenSlot;
vector<Matrix> potential;

Equation get_neighbor_equation(int idx_1, int idx_2){
    Equation res;
    res.constant_term = invN;
    res.terms.push_back(EquationTerm{ idx_1, idx_2, m32(-1) });

    m32 num_open_idx = ((N-1) - idx_1 - idx_2);

    // move an edge of type 1
    if(idx_1 != 0){
        // in -> out
        res.add(EquationTerm{ idx_1-1, idx_2,
            m32(1) 
                * invWeight * idx_1
                * (m32(1) - invNumOpenSlot * (num_open_idx + 1))
        });
        // in -> in
        res.add(EquationTerm{ idx_1, idx_2,
            m32(1) 
                * invWeight * idx_1
                * invNumOpenSlot * (num_open_idx + 1)
        });
    }
    if(idx_1 != num_edge1){
        // out -> out
        res.add(EquationTerm{ idx_1, idx_2,
            m32(1) 
                * invWeight * (num_edge1-idx_1)
                * (m32(1) - invNumOpenSlot * num_open_idx)
        });
        // out -> in
        res.add(EquationTerm{ idx_1+1, idx_2,
            m32(1) 
                * invWeight * (num_edge1-idx_1)
                * invNumOpenSlot * num_open_idx
        });
    }

    // move an edge of type 2
    if(idx_2 != 0){
        // in -> out
        res.terms.push_back(EquationTerm{ idx_1, idx_2-1,
            m32(1) 
                * invWeight * idx_2 * 2
                * (m32(1) - invNumOpenSlot * (num_open_idx + 1))
        });
        // in -> in
        res.add(EquationTerm{ idx_1, idx_2,
            m32(1) 
                * invWeight * idx_2 * 2
                * invNumOpenSlot * (num_open_idx + 1)
        });
    }
    if(idx_2 != num_edge2){
        // out -> out
        res.add(EquationTerm{ idx_1, idx_2,
            m32(1) 
                * invWeight * (num_edge2-idx_2) * 2
                * (m32(1) - invNumOpenSlot * num_open_idx)
        });
        // out -> in
        res.add(EquationTerm{ idx_1, idx_2+1,
            m32(1) 
                * invWeight * (num_edge2-idx_2) * 2
                * invNumOpenSlot * num_open_idx
        });
    }
    for(auto& term : res.terms) term.coeff /= (invWeight * invNumOpenSlot);
    res.constant_term /= (invWeight * invNumOpenSlot);
    //res.output();
    return res;
}


int main() {
    cin >> N;
    invN = m32(N).inv();

    num_edge1 = num_edge2 = 0;
    initial_state.assign(N, make_pair(0,0));
    rep(i,N-1){
        int u,v,c; cin >> u >> v >> c; u--; v--;
        if(c == 1){
            num_edge1++;
            initial_state[u].first++;
            initial_state[v].first++;
        }
        else{
            num_edge2++;
            initial_state[u].second++;
            initial_state[v].second++;
        }
    }
    if(num_edge1 == 0){
        swap(num_edge1, num_edge2);
        for(auto& a : initial_state) swap(a.first, a.second);
    }

    invWeight = m32(num_edge1 + 2 * num_edge2).inv();
    invNumOpenSlot = m32(N * (N-1) / 2 - (N-2)).inv();

    potential.assign(num_edge1+1, Matrix(num_edge2+1, num_edge1+2));
    const int constant_term_idx = num_edge1 + 1;
    const int equation_size = constant_term_idx + 1;
    for(int idx_1 = 0; idx_1 <= num_edge1; idx_1++) potential[idx_1][0][idx_1] = 1;

    for(int idx_2 = 1; idx_2 <= num_edge2; idx_2++){
        for(int idx_1 = 0; idx_1 <= num_edge1; idx_1++){
            auto eq = get_neighbor_equation(idx_1, idx_2-1);
            auto pcoeff = -eq.get_coeff(idx_1, idx_2);
            eq.add(EquationTerm{ idx_1, idx_2, pcoeff });
            pcoeff = pcoeff.inv();
            for(auto eqterm : eq.terms){
                rep(i,equation_size) potential[idx_1][idx_2][i] += potential[eqterm.idx_1][eqterm.idx_2][i] * eqterm.coeff;
            }
            potential[idx_1][idx_2][constant_term_idx] += eq.constant_term;
            rep(i,equation_size) potential[idx_1][idx_2][i] *= pcoeff;
        }
    }

    Matrix final_equation = Matrix(num_edge1+1, equation_size);
    for(int idx_1 = 0; idx_1 < num_edge1; idx_1++){
        auto eq = get_neighbor_equation(idx_1, num_edge2);
        for(auto eqterm : eq.terms){
            rep(i,equation_size) final_equation[idx_1][i] += eqterm.coeff * potential[eqterm.idx_1][eqterm.idx_2][i];
        }
        final_equation[idx_1][constant_term_idx] += eq.constant_term;
    }
    rep(i,equation_size) final_equation[num_edge1][i] += potential[num_edge1][num_edge2][i];
    rep(i,equation_size) final_equation[num_edge1][i] += potential[1][0][i] * num_edge1;
    rep(i,equation_size) final_equation[num_edge1][i] += potential[0][1][i] * num_edge2;

    for(int idx_1 = 0; idx_1 <= num_edge1; idx_1++) final_equation[idx_1][constant_term_idx] *= -1;

    auto potential_solution = final_equation.linear_equation().first;
    potential_solution.push_back(1);

    m32 ans = 0;
    for(auto [idx_1, idx_2] : initial_state){
        //cout << idx_1 << ", " << idx_2 << endl;
        rep(i,equation_size) ans += potential_solution[i] * potential[idx_1][idx_2][i];
    }
    cout << ans.val() << endl;
    return 0;
}


struct ios_do_not_sync{
    ios_do_not_sync(){
        ios::sync_with_stdio(false);
        cin.tie(nullptr);
    }
} ios_do_not_sync_instance;
yukicoder

結果

ソースコード
