#include 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(Clock::now()-start).count(); } }; struct Model { struct Node { uint16_t v, a, b; }; int N, K, T; vector cell, id; vector dry, x, q, p, adj, dq; vector order; // 各落雷位置について T 個、先頭は落雷位置自身 Model(int n, int k, const vector& board): N(n), K(k), id(n*n, -1) { for (int y=0; y=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; selapsed()>=limit) { order.clear(); return false; } int sy=cell[s]/N, sx=cell[s]%N; vector ys, xs; for (int y=sy; y>=0; --y) ys.push_back(y); for (int y=sy+1; y=0; --z) xs.push_back(z); for (int z=sx+1; 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& z, double penalty, vector& grad, double& loss, double& expected, double& sum_x, const Timer* timer=nullptr, double limit=0) { sum_x=0; for (int v=0; velapsed()>=limit) return false; const Node* nodes=order.data()+size_t(s)*T; p[s]=q[s]; expected+=p[s]; for (int i=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 m,v; double beta1_pow=1, beta2_pow=1; explicit Adam(int n): m(n),v(n) {} void step(vector& z, const vector& 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 board(N); for (auto& s:board) cin>>s; Model model(N,K,board); mt19937_64 rng(Config::SEED); uniform_real_distribution uniform(-1.0,1.0); vector 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) break; adam.step(z,grad); ++iterations; } } // sigmoid は単調なので、最終 z の上位 K 個と最終 x の上位 K 個は同じ。 vector 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