Submission #1205128

#	Time	Username	Problem	Language	Result	Execution time	Memory
1205128		phamtienhoang	Museum (CEOI17_museum)	C++17	80 / 100	454 ms	1114112 KiB

museum

#include <iostream>
#include <algorithm>
#include <cstring>
using namespace std;
using ll = long long;

const ll INF = (ll)1e18;
const int MAXN = 10005;

// input
int n, k, root;
// adjacency via edge-list
int head[MAXN], nxt[2*MAXN], to[2*MAXN], wgt[2*MAXN], ec=0;
void addEdge(int u, int v, int w){
  nxt[++ec] = head[u];
  head[u] = ec;
  to[ec] = v;
  wgt[ec] = w;
}

// DP: 
//   f[u][j][0] = best *closed* cost to pick j rooms in u’s subtree, start-and-end at u  
//   f[u][j][1] = best *open*   cost to pick j rooms, start at u, end anywhere  
ll f[MAXN][MAXN][2];
// scratch arrays for merging children: curr[state][j], nextdp[state][j]
ll curr[2][MAXN], nextdp[2][MAXN];
int sz[MAXN];

void dfs(int u, int p){
  // recurse children first
  for(int e = head[u]; e; e = nxt[e]){
    int v = to[e];
    if(v==p) continue;
    dfs(v, u);
  }

  // ---- FIXED INITIALIZATION ----
  // Before merging any child, we have only picked u itself (j=1), cost=0
  for(int s0 = 0; s0 < 2; s0++){
    for(int j = 0; j <= k; j++){
      curr[s0][j] = INF;
    }
  }
  curr[0][1] = 0;    // closed, 1 room
  curr[1][1] = 0;    // open,   1 room
  sz[u] = 1;
  // -----------------------------

  // now merge each child's DP
  for(int e = head[u]; e; e = nxt[e]){
    int v = to[e], c = wgt[e];
    if(v==p) continue;
    int su = sz[u], sv = sz[v];
    int ns = min(k, su + sv);

    // clear nextdp
    for(int s0 = 0; s0 < 2; s0++){
      for(int j = 0; j <= ns; j++){
        nextdp[s0][j] = INF;
      }
    }

    // merge child v
    for(int i = 1; i <= su; i++){
      ll cu0 = curr[0][i];
      ll cu1 = curr[1][i];
      if(cu0>=INF && cu1>=INF) continue;

      // skip child entirely (j2=0 ⇒ cost 0)
      nextdp[0][i] = min(nextdp[0][i], cu0);
      nextdp[1][i] = min(nextdp[1][i], cu1);

      // take j2=1..sv rooms from child
      for(int j2 = 1; j2 <= sv && i+j2 <= k; j2++){
        ll cl = f[v][j2][0];
        ll op = f[v][j2][1];
        if(cl < INF){
          ll cost_cl = cl + 2LL*c;      // go & return
          nextdp[0][i+j2] = min(nextdp[0][i+j2], cu0 + cost_cl);
          nextdp[1][i+j2] = min(nextdp[1][i+j2], cu1 + cost_cl);
        }
        if(op < INF){
          ll cost_op = op + 1LL*c;      // go but not return
          nextdp[1][i+j2] = min(nextdp[1][i+j2], cu0 + cost_op);
        }
      }
    }

    // copy back and update size
    sz[u] = ns;
    for(int s0 = 0; s0 < 2; s0++){
      for(int j = 1; j <= sz[u]; j++){
        curr[s0][j] = nextdp[s0][j];
      }
    }
  }

  // write into f[u]
  for(int j = 1; j <= sz[u]; j++){
    f[u][j][0] = curr[0][j];
    f[u][j][1] = curr[1][j];
  }
}

int main(){
  ios::sync_with_stdio(false);
  cin.tie(nullptr);

  cin >> n >> k >> root;
  memset(head, 0, sizeof(head));
  // read edges
  for(int i = 1; i < n; i++){
    int u,v,w;
    cin >> u >> v >> w;
    addEdge(u,v,w);
    addEdge(v,u,w);
  }
  // initialize f to INF
  for(int u = 1; u <= n; u++){
    for(int j = 0; j <= k; j++){
      f[u][j][0] = f[u][j][1] = INF;
    }
  }

  dfs(root, 0);
  // print the *open* tour cost for exactly k rooms
  cout << f[root][k][1] << "\n";
  return 0;
}

• Subtask 1: 20 / 20
• Subtask 2: 25 / 25
• Subtask 3: 35 / 35
• Subtask 4: 0 / 20