日々drdrする人のメモ

今日も明日もdrdr

CodeChef July Challenge 2019: Hit the Coconuts

CodeChef July Challenge 2019 の問題: Hit the Coconuts (Code: CCC)
問題ページ: Contest Page | CodeChef

問題概要

 N個のココナッツがあり、 i番目のココナッツは A_i回叩くと割れる。
しかし、 N個のココナッツはシャッフルされ、どのココナッツが何回叩けば割れるか分からない。

今回、 Z個のココナッツを割る必要がある。
少なくとも何回叩けば Z個のココナッツを割ることができるか求めよ。

制約

  • 1つのテストケースに含まれるケース数:  1 \le T \le 1000
  •  1 \le Z \le N \le 10^3
  •  1 \le A_i \le 10^6
  • ケース通して  N \cdot Z の和は  3 \times 10^6 を超えない

解法

この問題では、少なくとも何回叩けばよいかを知りたいので、最悪ケースの叩く順序を考えていく。

ここでは  A_1 \le A_2 \le ... \le A_{N-1} \le A_{N} とする。

例えば、 i番目のココナッツを割りたい場合、目的のココナッツが割れるまで A_i回ずつ叩いていくと少ない回数で達成できる。
最悪ケースでは、目的のココナッツを割るまでに  i \le j の全てのココナッツを余計に叩くことになる。つまり、 (N - i + 1) \times A_i 回叩くことになる。

f:id:smijake3:20190715232551p:plain
i番目を割ろうとする時の最悪ケースのイメージ
赤が目的のココナッツを叩く箇所、黄色が余計に叩く箇所)

この回数が最小となるココナッツを1個選択すれば部分点が取れる。(計算量:  O(N \log N))



次に複数のココナッツを割ることを考えるが、回数が小さい方のココナッツから割っていくことにする。(大きい方から割ることを考えてもほぼ同じ方針になる)

直前に i番目のココナッツを叩き割ったとする。次に j (\ge i)番目のココナッツを割ることを考える。

この時、 j番目のココナッツを叩き割るのに必要な回数は  (N - j + 1) \times (A_j - A_i) 回になる。
(この時、 A_i 回叩いて割れなかったココナッツは既知なので、それらのココナッツのみを叩けばよい)

f:id:smijake3:20190715232640p:plain
複数のココナッツを割る時のイメージ
(赤は目的のココナッツを叩く箇所、黄色は余計に叩く箇所)

これらを踏まえ、以下のDPを考える

 dp[i][z] =  i 番目までのココナッツを z個割った時の最小回数

このDPは以下の式で計算できる。

 dp[i][1] = (N + 1 - i) \times A_i

 {
\begin{align}
dp[i][z] & = \min_{1 \le j < i}(dp[j][z-1] + (N+1-i)(A_i - A_j)) \\
& = (N+1-i)A_i + \min_{1 \le j < i}(A_j \cdot i - (N+1) \cdot A_j + dp[j][z-1]) \\
& = (N+1-i)A_i + \min_{1 \le j < i}( f_j(i) ) \\
& (f_j(x) = A_j \cdot x - (N+1) \cdot A_j + dp[j][z-1])
\end{align}
}

このDPは、各zごとにConvex Hull Trickを用いて計算することができる。

そして、  \displaystyle \max_{1 \le i \le N} dp[i][Z] が解になる。

実装

部分点の提出コード(Python3): Solution: 25112502 | CodeChef


満点の実装は Li Chao Tree を用いた実装。
計算量は  O(Z N \log N)

提出コード (C++14): Solution: 25116098 | CodeChef

#define N 1003

class LiChaoTree {
  int n;
  ll xs[4*N];
  ll p[4*N], q[4*N];
  bool u[4*N];

  void _add_line(ll a, ll b, int k, int l, int r) {
    while(r-l > 0) {
      int m = (l + r) >> 1;
      if(!u[k]) {
        p[k] = a; q[k] = b;
        u[k] = true;
        return;
      }

      ll lx = xs[l], mx = xs[m], rx = xs[r-1];
      ll pk = p[k], qk = q[k];
      bool left = (a*lx+b < pk*lx+qk);
      bool mid = (a*mx+b < pk*mx+qk);
      bool right = (a*rx+b < pk*rx+qk);
      if(left && right) {
        p[k] = a; q[k] = b;
        return;
      }
      if(!left && !right) {
        return;
      }
      if(mid) {
        swap(p[k], a);
        swap(q[k], b);
      }
      if(left != mid) {
        k = 2*k+1; r = m;
      } else {
        k = 2*k+2; l = m;
      }
    }
  }

  ll _query(int k, ll x) {
    k += n - 1;
    ll s = u[k] ? p[k]*x+q[k] : INFLL;
    while(k > 0) {
      k = (k - 1) / 2;
      if(u[k]) {
        ll r = p[k]*x+q[k];
        s = mind(s, r);
      }
    }
    return s;
  }

public:
  LiChaoTree(int n0, ll ps[N]) {
    n = 1;
    while(n < n0) n <<= 1;

    rep(i, 2*n) u[i] = false;
    rep(i, n0) xs[i] = ps[i];
    repl(i, n0, 2*n-1) xs[i] = INF;
  }

  void add_line(ll a, ll b) {
    _add_line(a, b, 0, 0, n);
  }

  void add_segment_line(ll a, ll b, int l, int r) {
    ll l0 = l + n, r0 = r + n;
    ll s0 = l, t0 = r, sz = 1;
    while(l0 < r0) {
      if(r0 & 1) {
        --r0; t0 -= sz;
        _add_line(a, b, r0-1, t0, t0+sz);
      }
      if(l0 & 1) {
        _add_line(a, b, l0-1, s0, s0+sz);
        ++l0; s0 += sz;
      }
      l0 >>= 1; r0 >>= 1;
      sz <<= 1;
    }
  }

  ll query(int i) {
    return _query(i, xs[i]);
  }
};

ll a[N], xs[N];
ll dp[N][N];
int main() {
  int t;
  ll n, z;
  rep(i, N) xs[i] = i;
  cin >> t;
  while(t--) {
    cin >> n >> z;
    rep(i, n) cin >> a[i];
    sort(a, a+n);

    // dp[z][i] := (0-indexedの) i番目のココナッツまでをz個割った時の最小回数

    // 初期値の計算
    rep(i, n) dp[1][i] = a[i] * (n - i);

    repl(k, 1, z-1) {
      LiChaoTree lct(n, xs);

      // i = 0...k-1 について dp[k+1][i] は解がないため
      // 直線を追加するだけ
      rep(i, k) {
        lct.add_line(a[i], - n*a[i] + dp[k][i]);
        dp[k+1][i] = 1e15;
      }

      // CHTを用いたDPの計算
      repl(i, k, n-1) {
        ll v = (n-i) * a[i] + lct.query(i);
        dp[k+1][i] = v;
        lct.add_line(a[i], - n*a[i] + dp[k][i]);
      }

    }
    ll ans = INFLL;
    repl(i, z-1, n-1) {
      ans = min(ans, dp[z][i]);
    }
    cout << ans << endl;
  }
  return 0;
}