AOJ 2095 : Nagashi Soumen

問題リンク : Nagashi Soumen | Aizu Online Judge

問題概要 :
(x,y,z)からなるN個の点があたえられる
最大K本の線をひけるのだが、N個の点全てがK本の線のいずれか１つに属するように線を引かなければならない
尚且つ２つの点p1,p2を結ぶ際に、p1.z > p2.z を満たさなければならない
線を引く際にその線の長さ分だけコストがかかる
条件を満たすように線を引いた際にかかるコストの最小値をもとめよ

解法 :
最小費用流
１つの点を2つに分ける
p -> IN(p)-OUT(p) とする
この間のコストは十分に小さい値-INFとする
あとは関係を満たすものに辺を追加し流せるだけフローを流す
その結果にN*-INFを加算すると答えとなる
N個の点は全て通らなければならないので、その点にとても小さい値を割り当てることで全ての点が選ばれるようにする

コード :

#include<bits/stdc++.h>
 
#define REP(i,s,n) for(int i=s;i<n;i++)
#define rep(i,n) REP(i,0,n)
#define EPS (1e-10)
#define equals(a,b) (fabs((a)-(b))<EPS)
 
using namespace std;

bool LT(double a,double b) { return !equals(a,b) && a < b; }
bool LTE(double a,double b) { return equals(a,b) || a < b; }
 
struct Point {
  double x,y,z;
  Point operator - ( const Point &p ) const { return (Point){x-p.x,y-p.y,z-p.z}; }
  bool operator < ( const Point &p ) const {
    return ( ( !equals(z,p.z) ) ? z > p.z : ( ( !equals(y,p.y) ) ? y < p.y : LT(x,p.x) ) );
  }
};

typedef pair<double,int> ii;

struct Edge{
  int to,cap; double cost; int rev;
  Edge(int to=0,int cap=0,double cost=0,int rev=0):to(to),cap(cap),cost(cost),rev(rev){}
};

const int MAX_V = 1100, IINF = INT_MAX;
int V;
vector<Edge> G[MAX_V];
double h[MAX_V],dist[MAX_V];
int prevv[MAX_V],preve[MAX_V];

inline void add_edge(int from,int to,int cap,double cost){
  G[from].push_back(Edge(to,cap,cost,G[to].size()));
  G[to].push_back(Edge(from,0,-cost,G[from].size()-1));
}

double min_cost_flow(int s,int t){
  double res = 0;
  fill(h,h+V,0);
  while(1){
    priority_queue<ii,vector<ii>,greater<ii> > Q;
    fill(dist,dist+V,IINF);
    dist[s] = 0;
    Q.push(ii(0,s));
    while(!Q.empty()){
      ii p = Q.top(); Q.pop();
      int v = p.second;
      if( dist[v] < p.first ) continue;
      for(int i=0;i<G[v].size();i++){
        Edge &e = G[v][i];
        if( LT(0,e.cap) && !equals(dist[v]+e.cost+h[v]-h[e.to],dist[e.to]) && dist[e.to] > dist[v] + e.cost + h[v] - h[e.to] ) { 
          dist[e.to] = dist[v] + e.cost + h[v] - h[e.to];
          prevv[e.to] = v;
          preve[e.to] = i;
          Q.push(ii(dist[e.to],e.to));
        }
      }
    }
    //if( dist[t] == IINF ) break;
    rep(v,V) h[v] += dist[v];
    if( h[t] >= 0LL ) break;
    int d = IINF;
    for(int v=t;v!=s;v=prevv[v]) d = min(d,G[prevv[v]][preve[v]].cap);
    res += d * h[t];
    for(int v=t;v!=s;v=prevv[v]){
      Edge &e = G[prevv[v]][preve[v]];
      e.cap -= d;
      G[v][e.rev].cap += d;
    }
  }
  return res;
}


int N,K;
Point ps[100];

const double MAX = 10000;

double getDist(Point p){ return sqrt(p.x*p.x+p.y*p.y+p.z*p.z); }

void compute(){
  if( N <= K ) { puts("0"); return; }
  map<int,int> counter;
  bool failed = false;
  rep(i,N) counter[(int)ps[i].z]++;
  for(map<int,int>::iterator it = counter.begin(); it != counter.end(); it++ ) {
    if( it->second > K ) { failed = true; break; }
  }
  if( failed ) { puts("-1"); return; }
  rep(i,N*4) G[i].clear();
  int source = N * 2;
  int sink = source + 2;
  V = sink +1;
  add_edge(source,source+1,K,0);
  rep(i,N) { 
    add_edge(source+1,i,1,0);
    add_edge(i,N+i,1,-MAX);
    add_edge(N+i,sink,1,0); 
  }

  rep(i,N) rep(j,N) if( i != j && ps[i].z > ps[j].z ){
    add_edge(N+i,j,1,getDist(ps[i]-ps[j]));
  }

  printf("%.10f\n",min_cost_flow(source,sink)+N*MAX);
}

int main(){
  while( cin >> N >> K, N|K ){
    rep(i,N) cin >> ps[i].x >> ps[i].y >> ps[i].z;
    compute();
  }
  return 0;
}

ちなみにJAGの解説だと探索となっている
ただ、AOJだとMLEで通らない ( はず。
うまい実装すれば通るのだとうか
hashで状態をもってもメモリの上限の2.5倍くらいはかかった

一応そのコードも載せておく ( 答えは一致し、時間的にも大丈夫
コード :

const int MAX_V = 110;////////////////////////////////////////////////////////
const float DINF = 1e10;
int N,K;
Point ps[MAX_V];
//double memo[MAX_V][MAX_V][MAX_V][MAX_V];
typedef unsigned long long ull;
const ull base = 1000000007ULL;
unordered_map<ull,float> memo;
float mini;
 
float abs(Point p) { return sqrt(p.x*p.x+p.y*p.y+p.z*p.z); }
 
float rec(int cur,int *prev,int kept){
  if( cur >= N ) return 0;
  //rep(i,4) cout << prev[i] << " "; cout << endl;
  int st[4];
  rep(i,4) st[i] = prev[i];
  sort(st,st+4);
  ull key = 0ULL;
  rep(i,4) {
    key *= base;
    key += (ull)(st[i]+1);
  }
  if( memo.count(key) ) return memo[key];
  //double &ret = memo[st[0]+1][st[1]+1][st[2]+1][st[3]+1];
  //if( !equals(ret,DINF) ) return ret;
  float ret = DINF;
  rep(i,4){
    if( i && prev[i-1] == -1 ) continue;
    int temp = prev[i];
    if( prev[i] == -1 ) {
      if( kept + 1 > K ) continue;
      if( i && prev[i-1] > cur ) continue;
      prev[i] = cur;
      ret = min(ret,rec(cur+1,prev,kept+1));
    } else {
      if( LT(ps[cur].z,ps[prev[i]].z) ) {
        if( i && prev[i-1] > cur ) continue;
        prev[i] = cur;
        ret = min(ret,rec(cur+1,prev,kept)+abs(ps[cur]-ps[temp]));
      }
    }
    prev[i] = temp;
  }
  return memo[key] = ret;
}
 
int main(){
  while( cin >> N >> K, N|K ){
    map<int,int> counter;
    bool failed = false;
    rep(i,N) {
      cin >> ps[i].x >> ps[i].y >> ps[i].z;
      int cnt = ++counter[ps[i].z];
      if( cnt > K ) failed = true;
    }
    if( failed ) { puts("-1"); continue; }
    sort(ps,ps+N);
    //rep(i,N+2) rep(j,N+2) rep(k,N+2) rep(l,N+2) memo[i][j][k][l] = DINF;
    memo.clear();
    int temp[4] = {-1,-1,-1,-1};
    mini = rec(0,temp,0);
    if( equals(mini,DINF) ) puts("-1");
    else printf("%.10lf\n",mini);
  }
  return 0;
}