F - Fresh tomatoes

キーワード

最小全域木(最大全域木)

解説

$N$ 個のトマトに対応する $N$ 個の頂点からなる，辺の無いグラフを考えます。
そして，Ajinoko君の行動に，グラフへの操作を対応させます。

1. 箱からトマトを $1$ つ選び，袋に移す。選んだトマトに対応する頂点を $v$ とする。
2. 箱と袋からそれぞれ $1$ つずつトマトを選び，$1 \leq m \leq M$ を満たす整数 $m$ を選ぶ。箱と袋から選んだトマトに対応する頂点の番号をそれぞれ $i,j$ とし，新鮮度をそれぞれ $x,y$ とすると，重みが $(x+y)^m\pmod M$ の辺 $\{i,j\}$ を張る。

すると，全操作終了後のグラフは頂点 $v$ を根とする全域木となります。

なぜならば，箱に入っているトマトに対応する頂点の集合を $S$ ，袋に入っているトマトに対応する頂点の集合を $T$ とすると，行動$2$の操作を行う度に $S$ から頂点が $1$ つ削除され $T$ に追加されます。
このとき削除・追加された頂点を $i$ とすると，$i$ の親を頂点 $j \in T$ とする辺が張られます。行動 $2$ の操作は $S= \emptyset$ となるまで繰り返されるため，全操作終了後には $T=\{1,2,3,...,N\}$ ，かつ，$v$ を除く $T$ の任意の要素には親が存在するので，全域木が構成されます。

全操作終了後の元気度は，全域木の重み(各辺の重みの総和)の最大値であるため，求める答えは最大全域木の重みになります。

最大全域木の重みは，プリム法やクラスカル法などによって求めることができます。
グラフの頂点の数を $V$ ，辺の数を $E$ とすると，どちらも計算量は $O(E \log V)$ です。
今回の問題では辺の数は $O(N^{2}M)$ なので，すべての辺を張ろうとすると計算量は $O(N^{2}M \log N)$ となり実行時間制限を超過する場合があります。

そこで少し工夫をします。
ここではプリム法とクラスカル法のどちらも重みが最大の辺から見ていくので，頂点 $i,j$ を端点とする辺を選ぶ際は重みが最大のものを選ぶことになります。
よって，頂点 $i,j$ を端点とする辺は，重みが最大の辺のみを張れば良いです。

$2 \leq x+y \leq 2000$ と範囲が小さいので，$2 \leq a \leq 2000$ を満たす任意の $a$ に対し $a^m \pmod M$ の最大値を求めておくことで，任意の異なる $2$ 頂点 $i,j$ に対し辺 $\{i,j\}$ の重みの最大値を $O(1)$ で取り出すことができます。

最終的な計算量は，必要な辺のみを張る部分で$O(M \max {A} +N^{2})$，その後のプリム法またはクラスカル法で$O(N^{2} \log N)$なので，全体で$O(M \max A +N^{2} \log N)$ です。

プリム法による実装例(C++)

xxxxxxxxxx

#include <bits/stdc++.h>

using namespace std;

​

struct UnionFind {

  vector<int> par;

​

  UnionFind(int sz) {

    par.resize(sz, -1);

  }

  int root(int pos) {

    if(par[pos] < 0) return pos;

    else return par[pos] = root(par[pos]);

  }

  bool unite(int u, int v) {

    u = root(u); v = root(v);

    if(u == v) return false;

    par[v] = u;

    return true;

  }

  bool same(int u, int v) {

    return root(u) == root(v);

  }

};

​

int main(){

  int n, m; cin >> n >> m;

  vector<int> a(n);

  for(int i = 0; i < n; i++) cin >> a[i];

​

  // 前計算

  vector<vector<int>> power(2001, vector<int>(m + 1));

  for(int i = 0; i <= 2000; i++) power[i][0] = 1;

  for(int i = 2; i <= 2000; i++) {

    for(int j = 1; j <= m; j++) {

      power[i][j] = power[i][j - 1] * i;

      power[i][j] %= m;

    }

  } 

  vector<int> mx(2001, -1);

  for(int i = 2; i <= 2000; i++){

    for(int j = 1; j <= m; j++){

      mx[i] = max(mx[i], power[i][j]);

    }

  }

​

  // プリム法

  priority_queue<tuple<int,int,int>> q;

  for(int i = 0; i < n; i++) {

    for(int j = 0; j < i; j++) {

      q.push({mx[a[i] + a[j]], i, j});

    }

  }

​

  int res = 0;

  UnionFind uf(n);

  int edge_cnt = 0;

  while(edge_cnt < n - 1) {

    auto [cost, u, v] = q.top(); q.pop();

    if(uf.same(u, v)) continue;

​

    res += cost;

    uf.unite(u, v);

    edge_cnt++;

  }

​

  cout << res << '\n';

​

  return 0;

}

類題