けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 246 G - Game on Tree 3 (黄色, 600 点)

めっちゃ面白い問題だった!

問題概要

頂点 0 を根とする頂点数  N の根付き木が与えられます。頂点 0 以外の頂点  v には数値  A_{v} が書かれています。今、頂点 0 にコマが置いてあります。

高橋くんと青木くんが次の動作を交互に繰り返します。

  • 青木くんは、好きな頂点を選んで、その頂点の値を 0 にする
  • 高橋くんは、コマを子頂点のいずれかに移動させる。移動できないときはその場でゲームを終了する

高橋くんは任意のタイミングでゲームを終了できます。そしてゲーム終了した時のコマの置かれた頂点の数値が、ゲームのスコアとなります。

高橋くんはスコアの最大化を目指し、青木くんはスコアの最小化を目指します。双方が最善を尽くしたときのゲームのスコアを求めてください。

制約

  •  2 \le N \le 10^{5}

考えたこと

まず真っ先に二分探索したいと思った。そうすれば、白黒ゲームになる。つまり、各頂点が白または黒に塗られていて、

  • 高橋くんは黒頂点に着地することを目指す
  • 青木くんは黒頂点を白く塗りつぶすことで、高橋くんが黒頂点に着地するのを阻止することを目指す

というゲームになる。このように、二分探索することで、ゲームだったりグラフだったり盤面だったりを白黒にするケースはよくある気がする。

白黒ゲームの解き方

まず思ったのは、高橋くんが勝てる場合、勝つ直前は下図のように「コマのある頂点の子頂点のうち、2 個以上が黒色」という状況にある。

1 個だけが黒色だと、先に青木くんに白色にされて勝てない。でも 2 個が黒色なら青木くんが両方を白色にする前に高橋くんが黒色頂点に着地できるのだ。

そしてさらに言えば、3 個が黒色なら、コマがもう一歩手前にあっても間に合う。

この状況を一般化して、各頂点の「余裕値」を求められそうだ。つまり次のような値を計算できそうだ。


dp[v] ← 頂点  v にコマがある状態から、青木くんに何手先に操作されても、高橋くんが黒頂点に着地できるか


たとえば、

  • 黒頂点自体の余裕値は 1 (すでに着地していることから)
  • 子頂点のうちの 2 個が黒頂点であるような頂点の余裕値は 1
  • 子頂点のうちの 3 個が黒頂点であるような頂点の余裕値は 2

となる。もっと一般に、次のような木 DP で余裕値が求められる。

dp[v] =  \max(0, \sum_{c : v の子頂点} dp[c]  - 1) + (頂点  v が黒頂点 ? 1 : 0)

そして dp[0] が 1 以上なら高橋くんの勝ち、0 なら青木くんの勝ち。この DP に要する計算量は  O(N) となる。

よって二分探索法全体を合わせると、全体の計算量は  A = \max_{v} A_{v} として、 O(N \log A) と評価できる。

コード

#include <bits/stdc++.h>
using namespace std;
using Graph = vector<vector<int>>;

int main() {
    int N;
    cin >> N;
    vector<int> A(N, 0);
    for (int i = 1; i < N; ++i) cin >> A[i];
    Graph G(N);
    for (int i = 0; i < N-1; ++i) {
        int u, v;
        cin >> u >> v;
        --u, --v;
        G[u].push_back(v);
        G[v].push_back(u);
    }

    long long low = -1, high = 1LL<<60;
    while (high - low > 1) {
        long long x = (low + high) / 2;

        auto dfs = [&](auto self, int v, int p = -1) -> int {
            int res = 0;
            for (auto ch : G[v]) {
                if (ch == p) continue;
                res += self(self, ch, v);
            }
            return max(0, res - 1) + (A[v] >= x);
        };

        if (dfs(dfs, 0) >= 1) low = x;
        else high = x;
    }
    cout << low << endl;
}

別解

二分探索ではない方法で  O(N \log N) で求められる方法がある。さっきの DP で言えば、

dp[v][x] ← 数値が  x 以上である頂点を黒頂点であるとした場合の、頂点  v の余裕値

とする。配列 dp[v] の値は、 x に応じて単調減少となるが、変化する部分は高々  O(N) 箇所しかないのだ。

よって dp[v] の値は、変化点も含めて、高々  O(N) 次元ベクトルで管理できる。

木 DP を組むときに、マージテクを使い、さらにヒープのマージを高速にできる meldable heap を用いることで  O(N \log N) の計算量となる。

詳細は Nachia さんのユーザ解説にて

atcoder.jp

AtCoder ABC 246 F - typewriter (青色, 500 点)

包除原理を学べる問題!

問題概要

 N 個の文字列  S_{1}, S_{2}, \dots, S_{N} が与えられます。次の手順によって作れる長さ  L の文字列の個数を 998244353 で割ったあまりを求めてください。

  1.  k = 1, 2, \dots, N のいずれかを選ぶ
  2. 文字列  S_{k} に含まれる文字のみを使って、長さ  L の文字列を作る

制約

  •  1 \le N \le 18
  •  1 \le L \le 10^{9}

包除原理へ

たとえば  N = 2 L = 2 で文字列が "ab"、"ac" である場合は

  • "ab" から作れるのは、"aa", "ab", "ba", "bb" の 4 通り
  • "ac" から作れるのは、"aa", "ac", "ca", "cc" の 4 通り
  • 両方から作れるのは、"aa" の 1 通り

ということで、4 + 4 - 1 = 7 通りと求められます。より一般に  N = 2 の場合は、

(S[0] から作れる文字列の個数)
+ (S[1] から作れる文字列の個数)
- (S[0], S[1] の両方から作れる文字列の個数)

を計算すればよいことになります。

なお、S[0] から作れる文字列の個数は、S[0] に含まれる文字の種類数を  m として  m^{L} 通りとなります。

S[0], S[1] から作れる文字列の個数は、S[0] と S[1] にともに含まれる文字の種類数を  m として  m^{L} 通りとなります。

一般の  N の場合

 N = 2 の場合を拡張して、一般の  N の場合は次のように求められます。


文字列  S_{1}, S_{2}, \dots, S_{N} の空でない部分集合  T 2^{N} - 1 通り考えられる。そのそれぞれに対して、

  • 集合  T に含まれる文字列のいずれからも、作れる長さ  L の文字列の個数を  f(T) としたとき
    •  T の要素数が奇数ならば  f(T) を加算
    •  T の要素数が偶数ならば  f(T) を減算

なお、 f(T) の計算は、 T に含まれる文字列のすべてに含まれる文字の個数を  m として、

 f(T) = m^{L}

と計算できます。

コード

計算量は、各部分集合  T に対して

  • すべてに含まれる文字の個数  m の計算: O(N)
  •  m^{L} の計算: O(\log L)

ですので、全体として  O(2^{N}(N + \log L)) となります。

#include <bits/stdc++.h>
using namespace std;

const int MOD = 998244353;
long long modpow(long long a, long long n, long long mod) {
    long long res = 1;
    while (n > 0) {
        if (n & 1) res = res * a % mod;
        a = a * a % mod;
        n >>= 1;
    }
    return res;
}

int main() {
    int N, L;
    cin >> N >> L;
    vector<int> S(N, 0);
    for (int i = 0; i < N; ++i) {
        string str;
        cin >> str;
        for (char c : str) S[i] |= 1 << (c - 'a');
    }

    long long res = 0;
    for (int bit = 1; bit < (1 << N); ++bit) {
        int tmp = (1 << 26) - 1;
        for (int i = 0; i < N; ++i) {
            if (bit & (1 << i)) tmp &= S[i];
        }
        int m = __builtin_popcount(tmp);
        long long mL = modpow(m, L, MOD);

        if (__builtin_popcount(bit) & 1) res = (res + mL) % MOD;
        else res = (res + MOD - mL) % MOD;
    }
    cout << res << endl;
}

AtCoder ABC 246 E - Bishop 2 (水色, 500 点)

迷路の最短路問題なので BFS でやりたくなるが、まともにやると  O(N^{3}) で TLE してしまう!!

頂点の持ち方を工夫して 0-1 BFS で解く!
別解として枝刈り BFS も。

drken1215.hatenablog.com

問題概要

 N \times N のグリッドが与えられます。各マスは通路 (文字 . で表される) と壁 (文字 # で表される) のいずれかです。

今あなたは座標 (Ax, Ay) にいます。ここから「ビショップの移動」を繰り返して座標 (Bx, By) に到達したいとします。そのための最小手数を答えてください。到達不可能な場合は -1 と出力してください。

なおビショップの移動とは「斜め四方向に、壁にぶつからない限りはどこまでも移動できる」というものです。壁を通り抜けたり、盤面外に出たりすることはできません。

制約

  •  2 \le N \le 1500

BFS では間に合わない

「迷路の最短路」は BFS の典型問題です。ただし通常の「迷路の最短路」問題では、1 回で移動できるマスが「上下左右に隣接する 4 マス」しかないのに対して、今回の問題では 1 回で  O(N) 個のマスに移動できうそうです。

BFS の計算量はグラフの辺数  E に対して  O(E) です。1 回で移動できる箇所が  O(N) となると、グラフの辺数は  O(N^{3}) になります。よって単なる BFS だと、計算量は  O(N^{3}) となって間に合いません。

 

ノードに「移動方向」を持たせて 1 マスずつ移動する

そこで、「1 回の移動で斜めにどこまでも進める」というのを次のように読み替えることにします。


  • 各頂点が (縦方向座標, 横方向座標, 前回の移動方向) を表すグラフを考える
  • 各頂点に対して斜め四方向に進んだ 4 マスにのみ辺を張る
    • 移動方向が前回のものと一致するときは「辺のコスト:0」
    • 移動方向が前回のものと異なるときは「辺のコスト:1」

こうしてできるグラフ上で、(Ax, Ay, 任意) を始点、(Bx, By, 任意) を終点とした最短経路問題を解けばよいことになります!!

こうしてグラフの辺数が  O(N^{2}) で抑えられることになりました。ただし元々のグラフでは辺のコストがすべて 1 だとみなせたのに対して、今回は辺のコストは 0 と 1 がありうることになります。

普通の BFS ではダメなので、Dijkstra 法や 0-1 BFS を使います。

 

0-1 BFS

0-1 BFS は、辺のコストが 0, 1 のみに限ることを利用して、Dijkstra 法における priority_queue の部分を deque にすることで計算量を落とす手法です。Dijkstra 法の計算量  O(V + E \log V) が、0-1 BFS では  O(V + E) に落ちます。

辺の重みが 0 と 1 のみに限られるとき、Dijkstra 法を回したときの priority_queue の中身は常に「最大値と最小値との差が 1 以下」という状態が保たれます。  

よって、priority_queue の中身は高々 2 種類の値しか同時に共存しないことになります。priority_queue のような高級なデータ構造を使う必要はありません。代わりに


両端キュー deque を使う


という方法が考えられます。そして辺  e = (u, v) について緩和するとき、

  • 辺のコストが 0 の場合:deque の先頭に push
  • 辺のコストが 1 の場合:deque の末尾に push

というようにすればよいことになります。deque の両端に push する操作は  O(1) でできるので、priority_queue を使うのに比べて計算量が落ちます。

今回の場合、計算量は  O(N^{2}) になります。

 

コード

#include <bits/stdc++.h>
using namespace std;

const int INF = 1 << 29;
const vector<int> dx = {1, 1, -1, -1};
const vector<int> dy = {1, -1, 1, -1};

template<class T> inline bool chmin(T& a, T b) {
    if (a > b) { 
        a = b; 
        return true; 
    } 
    return false;
}

int main() {
    // 入力
    int N, Ax, Ay, Bx, By;
    cin >> N >> Ax >> Ay >> Bx >> By;
    --Ax, --Ay, --Bx, --By;
    vector<string> S(N);
    for (int i = 0; i < N; ++i) cin >> S[i];

    // 頂点数は (N x N x 4)
    vector<vector<vector<int>>> dp(N, vector<vector<int>>(N, vector<int>(4, INF)));
    deque<tuple<int,int,int>> que;

    // 初期条件
    for (int dir = 0; dir < 4; ++dir) {
        dp[Ax][Ay][dir] = 1;  // 最初の 1 手をノーコストにするため
        que.push_back({Ax, Ay, dir});
    }

    // 0-1 BFS
    while (!que.empty()) {
        auto [x, y, dir] = que.front();
        que.pop_front();
        
        for (int ndir = 0; ndir < 4; ++ndir) {
            int nx = x + dx[ndir], ny = y + dy[ndir];
            if (nx < 0 || nx >= N || ny < 0 || ny >= N) continue;
            if (S[nx][ny] == '#') continue;

            int cost = (ndir != dir);
            if (chmin(dp[nx][ny][ndir], dp[x][y][dir] + cost)) {
                if (cost == 0) que.push_front({nx, ny, ndir});
                else que.push_back({nx, ny, ndir});
            }
        }
    }
    int res = INF;
    for (int dir = 0; dir < 4; ++dir) chmin(res, dp[Bx][By][dir]);
    cout << (res < INF ? res : -1) << endl;
}

別解:枝刈り BFS

さきほど、通常の BFS では TLE してしまうと書いたが、少し工夫することで通すことができる。

通常の BFS では、今いるマスから、斜め四方に壁にぶつからない限りどこまでも行く遷移を試す。

下のコードのような感じ。このコードは、移動方向 dir を決めた上で、その方向に壁にぶつかるまで遷移する様子を表している。

while (true) {
      x += dx[dir], y += dy[dir];
      if (x < 0 || x >= N || y < 0 || y >= N) break;
      if (S[x][y] == '#') break;
                
      if (chmin(dp[x][y], cur_dist + 1)) {
             que.push({x, y});
      } 
}

枝刈り BFS は、次のように「もともとの距離値がより小さいようなマスに到達したら打ち切る」ようにするだけ!!!

while (true) {
      x += dx[dir], y += dy[dir];
      if (x < 0 || x >= N || y < 0 || y >= N) break;
      if (S[x][y] == '#') break;

      if (dp[x][y] < cur_dist + 1) break;
                
      if (chmin(dp[x][y], cur_dist + 1)) {
             que.push({x, y});
      } 
}

たったこれだけの工夫で劇的に早くなる。

コード

#include <bits/stdc++.h>
using namespace std;

const int INF = 1 << 29;
const vector<int> dx = {1, 1, -1, -1};
const vector<int> dy = {1, -1, 1, -1};

template<class T> inline bool chmin(T& a, T b) {
    if (a > b) { 
        a = b; 
        return true; 
    } 
    return false;
}

int main() {
    // 入力
    int N, Ax, Ay, Bx, By;
    cin >> N >> Ax >> Ay >> Bx >> By;
    --Ax, --Ay, --Bx, --By;
    vector<string> S(N);
    for (int i = 0; i < N; ++i) cin >> S[i];

    // 頂点数は (N x N)
    vector<vector<int>> dp(N, vector<int>(N, INF));
    queue<pair<int,int>> que;

    // 初期条件
    dp[Ax][Ay] = 0;
    que.push({Ax, Ay});

    // 枝刈り BFS
    while (!que.empty()) {
        auto [cx, cy] = que.front();
        que.pop();
        int cur_dist = dp[cx][cy];
        for (int dir = 0; dir < 4; ++dir) {
            int x = cx, y = cy;
            while (true) {
                x += dx[dir], y += dy[dir];
                if (x < 0 || x >= N || y < 0 || y >= N) break;
                if (S[x][y] == '#') break;

                if (dp[x][y] < cur_dist + 1) break;
                
                if (chmin(dp[x][y], cur_dist + 1)) {
                    que.push({x, y});
                } 
            }
        }
    }
    cout << (dp[Bx][By] < INF ? dp[Bx][By] : -1) << endl;
}

AtCoder ABC 246 D - 2-variable Function (緑色, 400 点)

前回の ABC に続いて、D 問題が数学っぽい。でも実際には今回は探索問題ですね。

問題概要

非負整数  a, b を用いて

 X = a^{3} + a^{2}b + a b^{2} + b^{3}

と表せる整数  X を「特別数」と呼ぶことにします。

非負整数  N が与えられますので、 N 以上である最小の「特別数」を求めてください。

制約

  •  0 \le N \le 10^{18}

変数が複数あるときは 1 つを固定する

見た目は確かに数学色の強い問題だけど、考えるべき基礎は決して特別なものではない。

このように  a, b という 2 つの変数が絡む問題で考えるべきことは、1 つの変数を固定して考えるというもの。

ここでは  a を固定してみよう。たとえば  a = 5 としたとしてみよう。このとき、

 X = b^{3} + 5 b^{2} + 25 b + 125

となる。これが  N 以上となるような最小の  b を求めたいと思ったら、それはまさに二分探索法がどんぴしゃだ!!

つまり次のように求められる。

// 式計算
long long func(long long a, long long b) {
    return a*a*a + a*a*b + a*b*b + b*b*b;
}

// a を固定したときの N 以上の最小の特別数を求める
long long solve(long long a) {
    long long low = -1, high = (十分大きな値);
    while (high - low > 1) {
        long long b = (low + high) / 2;
        if (func(a, b) >= N) high = b;
        else low = b;
    }
    return func(a, high);
}

変数の範囲を考える

最後に検討すべきことは以下の 2 点だ

  • 変数  a の調べるべき範囲
  • 変数  a を固定したとき、二分探索法の初期値の決め方

これを考える前に、 X = 10^{18} が「特別数」であることを確認しよう。具体的には  a = 10^{6},  b = 0 としたとき

 a^{3} + a^{2} b + a b^{2} + b^{3} = 10^{18}

となるのだ。よって  N \le 10^{18} という制約を考えると、 N 以上の最小の特別数が  10^{18} より大きくなる可能性は考えなくて良いのだ!!!

そのことから、 a \gt 10^{6} の範囲を調べる必要がないことが言える。なぜなら  a \gt 10^{6} とすると

 a^{3} + a^{2} b + a b^{2} + b^{3} \ge a^{3} \gt 10^{18}

となるためだ。 10^{18} より大きくなる範囲は考慮不要なので、 a \gt 10^{6} は考えなくてよい。

同様に  a を固定したとき、 a^{3} + a^{2} b + a b^{2} + b^{3} N 以上となる最小の  b を求めるための二分探索法の上限値も、 10^{6} より大きく取る必要はない!!

以上より、次の結論が言える


  • 変数  a 0, 1, 2, \dots, 10^{6} まで調べれば十分
  • 変数  a を固定したとき、二分探索法の上限値の初期値は  10^{6} とすればよい

計算時間的にも十分に間に合う。一応計算量をちゃんと書くと、 O(N^{\frac{1}{3}} \log N) と評価できる。

コード

#include <bits/stdc++.h>
using namespace std;

long long func(long long a, long long b) {
    return a*a*a + a*a*b + a*b*b + b*b*b;
}

long long solve(long long N, long long a) {
    long long low = -1, high = 1000000;
    while (high - low > 1) {
        long long b = (low + high) / 2;
        if (func(a, b) >= N) high = b;
        else low = b;
    }
    return func(a, high);
}

int main() {
    long long N;
    cin >> N;

    long long res = 1000000000000000000LL;  // 10^18 が上限値
    for (long long a = 0; a <= 1000000; ++a) {
        res = min(res, solve(N, a));
    }
    cout << res << endl;
}

 

別解:尺取り法

上述の「変数固定 + 二分探索」が比較的思いつきやすいと思う。一方そのような問題の多くは、尺取り法でも解けることが多いと思う。

そうすると計算量は  O(N^{\frac{1}{3}}) に落ちる。注目することは

  •  a を増やしていくとき、
  •  a^{3} + a^{2} b + a b^{2} + b^{3} N 以上となるような最小の  b は単調減少していく

ということだ。よって、 a を増やしながら、 b を減らしていくという「尺取り法」で十分高速に解ける。詳しくは公式解説へ!!

atcoder.jp

AtCoder ABC 246 C - Coupon (灰色, 300 点)

ソートが必要になるところが少し難しいかもしれない。

問題概要

 N 個の商品があって、それぞれ  A_{1}, A_{2}, \dots, A_{N} 円である。

またクーポンが  K 枚あって、1 枚のクーポンを使って次のことができる。

  • 商品を 1 つ選ぶ
  • その商品の価格を  X 円減少させる
  • ただしもとの価格が  X 円未満の場合は、0 円にする

クーポンを効率よく活用したとき、 N 個の商品をすべて購入するのに要する最小費用を求めよ。

制約

  •  1 \le N \le 2 \times 10^{5}
  •  1 \le K \le 10^{9}

手を動かしてみる

まずは手を動かして問題の様子をつかんでみよう。たとえば  A = (12, 16, 18) K = 10,  X = 5 としてみよう。

このとき、たとえばまず 1 個目の商品にクーポンを適用してみよう。そうすると 2 回適用することで価格は、

12 → 7 → 2

となる。このときさらにクーポンを適用するのは少しもったいない。なぜならその場合、

  • 1 回目のクーポンによる価格減少分は、12 → 7 で「-5」
  • 2 回目のクーポンによる価格減少分は、7 → 2 で「-5」
  • 3 回目のクーポンによる価格減少分は、2 → 0 で「-2」

ということになる。1, 2 回目は「-5」を達成したのに対して、3 回目は「-2」にとどまってしまう。

それをするくらいなら、2 個目の商品の 16 円にクーポンを適用した方がいい。そうすることで「-5」の達成を継続できる。

 

最後のあまりを処理

以上の話を一般化しよう。一般に「$-X$ 円」を達成できるところを優先的に減らしていこう。具体的には、各  A_{i} に対して、A[i] / X (あまりを切り捨てた値) の総和回だけ「$-X$ 円」を実施できることになる。

この総和を  S としたとき、もし  S \ge K ならば、 K 回のクーポンはすべて「$-X$ 円」を達成できるので、それが答えになる。

よって  S \lt K の場合を考えよう。この場合、まずは「$-X$ 円」を達成できる部分は使い切ってしまおう。そうすると

  • K -= S
  • i に対して A[i] %= X

という状態になったものと考えてよい。たとえば  A = (12, 16, 18),  K = 10,  X = 5 のとき、下図のように、 A = (2, 1, 3) K = 2 の問題に帰着される。

こうなったら、余ったクーポンを有効活用するために、「減らせる金額が大きいところを優先的に使っていく」のがよいと言える。具体的には次のように求められる。


  • 配列 A (あまりを取った値) を大きい順にソートして
  • 大きい順に  K 個を 0 にする (ただし  K \ge N の場合はすべて 0 にして終了)
  • 残った A の総和が答え

計算量は「ソート」に最も時間がかかるので、 O(N \log N)

 

コード

#include <bits/stdc++.h>
using namespace std;

long long solve() {
    // 入力 (sum は最初の総額)
    long long N, K, X, sum = 0;
    cin >> N >> K >> X;
    vector<long long> A(N);
    for (int i = 0; i < N; ++i) cin >> A[i], sum += A[i];

    // 「-X」を適用できるところを先に適用してしまう
    long long S = 0;
    vector<long long> amari(N);
    for (int i = 0; i < N; ++i) {
        S += A[i] / X;
        A[i] %= X;
    }

    // S >= K のときはすべてのクーポンを
    if (S >= K) return sum - K * X;

    // K > S のときは、減らせる金額が大きいところを優先的に
    K -= S;
    sort(A.begin(), A.end(), greater<long long>());
    for (int i = 0; i < min(N, K); ++i) A[i] = 0;
    long long res = 0;
    for (int i = 0; i < N; ++i) res += A[i];
    return res;
}

int main() {
    cout << solve() << endl;
}

類題情報

似た感じで解ける問題!

drken1215.hatenablog.com