#1195039 (C++23(gcc16)) No.5025 Forest Fire

• 問題: No.5025 Forest Fire
• ユーザー: twins_fuyu
• 提出日時: 2026-10-24 16:12:38
• 言語: C++23(gcc16) (gcc 16.2.0 + boost 1.92.0 + ACL)
• 結果: AC
• 実行時間: 1,890 ms / 2,000 ms
• コード長: 6,844 bytes
• コンパイル時間: 3,191 ms
• コンパイル使用メモリ: 372,360 KB
• 実行使用メモリ: 40,592 KB
• スコア: 67,131,716
• 最終ジャッジ日時: 2026-10-24 16:15:45
• 合計ジャッジ時間: 156,063 ms
• ジャッジサーバーID （参考情報）: judge4_0 / <nil>

ファイルパターン	結果
other	AC * 50

ソースコード

#include <bits/stdc++.h>
using namespace std;
using Clock = chrono::steady_clock;

namespace Config {
// TODO: 提出前に以下の2行を有効化して値を指定する。
// constexpr double TIME_LIMIT = ...; // 秒。入力・前計算を含む実行時間
// constexpr double PENALTY = ...;    // max(0, sum(x)-K)^2 の係数
// 未設定時は ./main TIME_LIMIT PENALTY と実行する（提出時は上を設定）。
constexpr double TIME_LIMIT = 1.9, PENALTY = 100.0;
// 上の TODO を有効化したら、上記 #ifndef 〜 #endif を削除する。
constexpr double LR = 0.001, BETA1 = 0.9, BETA2 = 0.999, EPS = 1e-8;
constexpr uint64_t SEED = 0; // 再現用乱数シード
}

struct Timer {
    Clock::time_point start = Clock::now();
    double elapsed() const { return chrono::duration<double>(Clock::now()-start).count(); }
};

struct Model {
    struct Node { uint16_t v, a, b; };
    int N, K, T;
    vector<int> cell, id;
    vector<double> dry, x, q, p, adj, dq;
    vector<Node> order; // 各落雷位置について T 個、先頭は落雷位置自身

    Model(int n, int k, const vector<string>& board): N(n), K(k), id(n*n, -1) {
        for (int y=0; y<N; ++y) for (int z=0; z<N; ++z) {
            if (board[y][z]=='0') continue;
            id[y*N+z] = (int)cell.size();
            cell.push_back(y*N+z);
            dry.push_back((board[y][z]-'0')*0.1);
        }
        T = (int)cell.size();
        if (T>=65535) throw runtime_error("too many trees");
        for (int& v:id) if (v<0) v=T; // T は常に燃焼確率 0 のダミー親
        x.resize(T); q.resize(T); dq.resize(T);
        p.resize(T+1); adj.resize(T+1);
    }

    // マンハッタン距離を減らす隣接マスが親。
    // 縦横それぞれ起点から外側へ並べると、その直積もトポロジカル順になる。
    bool build(const Timer* timer=nullptr, double limit=0) {
        order.reserve(size_t(T)*T);
        for (int s=0; s<T; ++s) {
            if (timer && timer->elapsed()>=limit) { order.clear(); return false; }
            int sy=cell[s]/N, sx=cell[s]%N;
            vector<int> ys, xs;
            for (int y=sy; y>=0; --y) ys.push_back(y);
            for (int y=sy+1; y<N; ++y) ys.push_back(y);
            for (int z=sx; z>=0; --z) xs.push_back(z);
            for (int z=sx+1; z<N; ++z) xs.push_back(z);
            for (int y:ys) for (int z:xs) {
                int v=id[y*N+z];
                if (v==T) continue;
                int a=T, b=T;
                if (y!=sy) a=id[(y+(y<sy ? 1 : -1))*N+z];
                if (z!=sx) b=id[y*N+z+(z<sx ? 1 : -1)];
                order.push_back({uint16_t(v), uint16_t(a), uint16_t(b)});
            }
        }
        return true;
    }

    static double sigmoid(double z) {
        if (z>=0) return 1.0/(1.0+exp(-z));
        double e=exp(z); return e/(1.0+e);
    }

    // L = E + penalty * max(0, sum(x)-K)^2
    // E は指定の独立近似モデルでの期待焼失本数（全 N^2 落雷位置の平均）。
    // 時間切れで途中終了した勾配は使わない。
    bool evaluate(const vector<double>& z, double penalty, vector<double>& grad,
                  double& loss, double& expected, double& sum_x,
                  const Timer* timer=nullptr, double limit=0) {
        sum_x=0;
        for (int v=0; v<T; ++v) {
            x[v]=sigmoid(z[v]); sum_x+=x[v]; q[v]=dry[v]*(1.0-x[v]);
        }
        fill(dq.begin(), dq.end(), 0.0);
        expected=0; p[T]=0;
        for (int s=0; s<T; ++s) {
            if (timer && timer->elapsed()>=limit) return false;
            const Node* nodes=order.data()+size_t(s)*T;
            p[s]=q[s]; expected+=p[s];
            for (int i=1; i<T; ++i) {
                auto [v,a,b]=nodes[i];
                double u=p[a]+p[b]-p[a]*p[b]; // P(A union B)
                p[v]=q[v]*u; expected+=p[v];
            }
            // 各 p[v] は目的関数 sum(p) に直接 1 回寄与する。
            fill(adj.begin(), adj.end(), 1.0);
            for (int i=T-1; i>=1; --i) {
                auto [v,a,b]=nodes[i];
                double g=adj[v], u=p[a]+p[b]-p[a]*p[b];
                dq[v]+=g*u;
                double h=g*q[v];
                adj[a]+=h*(1.0-p[b]); adj[b]+=h*(1.0-p[a]);
            }
            dq[s]+=adj[s]; // p[s]=q[s]
        }
        double inv=1.0/(N*N), excess=max(0.0, sum_x-K);
        expected*=inv; loss=expected+penalty*excess*excess;
        grad.resize(T);
        for (int v=0; v<T; ++v)
            grad[v]=(-dry[v]*dq[v]*inv+2.0*penalty*excess)*x[v]*(1.0-x[v]);
        return true;
    }
};

struct Adam {
    vector<double> m,v;
    double beta1_pow=1, beta2_pow=1;
    explicit Adam(int n): m(n),v(n) {}
    void step(vector<double>& z, const vector<double>& g) {
        beta1_pow*=Config::BETA1; beta2_pow*=Config::BETA2;
        for (int i=0; i<(int)z.size(); ++i) {
            m[i]=Config::BETA1*m[i]+(1.0-Config::BETA1)*g[i];
            v[i]=Config::BETA2*v[i]+(1.0-Config::BETA2)*g[i]*g[i];
            z[i]-=Config::LR*(m[i]/(1.0-beta1_pow))
                 /(sqrt(v[i]/(1.0-beta2_pow))+Config::EPS);
        }
    }
};

int main(int argc, char** argv) {
    Timer timer;
    ios::sync_with_stdio(false); cin.tie(nullptr);
    double limit=Config::TIME_LIMIT, penalty=Config::PENALTY;
    if (argc==3) { limit=stod(argv[1]); penalty=stod(argv[2]); }
    if (!isfinite(limit) || !isfinite(penalty) || limit<=0 || penalty<0) {
        cerr << "TODO: set Config::TIME_LIMIT and Config::PENALTY, "
                "or run ./main TIME_LIMIT PENALTY\n";
        return 1;
    }
    int N,K; if (!(cin>>N>>K)) return 0;
    vector<string> board(N); for (auto& s:board) cin>>s;
    Model model(N,K,board);
    mt19937_64 rng(Config::SEED);
    uniform_real_distribution<double> uniform(-1.0,1.0);
    vector<double> z(model.T), grad;
    for (double& v:z) v=uniform(rng);
    Adam adam(model.T);
    int iterations=0;
    double loss=0, expected=0, sum_x=0;
    if (model.T && model.build(&timer,limit)) {
        while (timer.elapsed()<limit) {
            if (!model.evaluate(z,penalty,grad,loss,expected,sum_x,&timer,limit)) break;
            if (timer.elapsed()>=limit) break;
            adam.step(z,grad); ++iterations;
        }
    }
    // sigmoid は単調なので、最終 z の上位 K 個と最終 x の上位 K 個は同じ。
    vector<int> selected(model.T); iota(selected.begin(),selected.end(),0);
    int count=min(K,model.T);
    partial_sort(selected.begin(),selected.begin()+count,selected.end(),[&](int a,int b) {
        return z[a]!=z[b] ? z[a]>z[b] : a<b;
    });
    cout<<count<<'\n';
    for (int i=0; i<count; ++i) {
        int pos=model.cell[selected[i]];
        cout<<pos/N<<' '<<pos%N<<'\n';
    }
    cerr<<"iterations="<<iterations<<" elapsed="<<timer.elapsed()<<'\n';
    return 0;
}