C++ の動作確認をしてみた（３７２）

C++の練習を兼ねて, AtCoder Regular Contest 108 の問題F (Paint Tree) を解いてみた.

■感想.
1. 問題F は, 解答方針が見えなかったので, 解説を参考に実装し, AC版まで到達出来たので良かったと思う.
2. 但し, 解答方針を絞り込む部分で, 非常に苦労した上, さらに, いくつか仮定を想定して, 実装する羽目になった.
※ ex. 良さとして存在しうる最小値を考慮(minBest), 累積和 0 の場合は, 2 の冪乗でカウントしないように集計(mPow2Sum), グラフが直線の場合の例外処理.
3. 実装が, スリムでないので, 正答者の方々の実装を勉強する必要があると感じた.
4. 時間を見つけて, 引き続き, 過去問を振り返っていきたいと思う.

本家のサイト AtCoder Regular Contest 108 解説の各リンクをご覧下さい.

■C++版プログラム(問題F／AC版).

// C++(GCC 9.2.1)
// 解き直し.
// https://atcoder.jp/contests/arc108/editorial/350
#include <bits/stdc++.h>
using namespace std;
using LL = long long;
using vi = vector<int>;
using vvi = vector<vi>;
using P = pair<int, int>;
#define repex(i, a, b, c) for(int i = a; i < b; i += c)
#define repx(i, a, b) repex(i, a, b, 1)
#define rep(i, n) repx(i, 0, n)
#define repr(i, a, b) for(int i = a; i >= b; i--)
#define pb push_back
const LL MOD = 1e9 + 7;
LL mPow2[202020]; // 2 の 冪乗(MOD版)

int main(){
    
    // 1. 入力情報.
    int N;
    scanf("%d", &N);
    vvi G(N);
    rep(i, N - 1){
        int a, b;
        scanf("%d %d", &a, &b);
        a--;
        b--;
        G[a].pb(b);
        G[b].pb(a);
    }
    
    // 2. 2 の 冪乗(MOD版) を 保存.
    mPow2[0] = 1;
    repx(i, 1, N) mPow2[i] = mPow2[i - 1] << 1, mPow2[i] %= MOD;
    
    // 3. 頂点 1 から 各頂点までの距離を計算.
    // https://ja.wikipedia.org/wiki/幅優先探索
    auto bfs = [&](vvi &G, int s, int* d){
        // 3-1. 空のキュー.
        queue<int> q;
        
        // 3-2. 探索地点 s をキュー q に追加.
        q.push(s);
        while(!q.empty()){
            // 3-3. キューから取り出す.
            int u = q.front();
            q.pop();
            
            // 3-4. 取り出した要素を処理.
            for(auto &e : G[u]){
                // 3-5. 訪問済であれば, 処理をスキップ.
                if(d[e])            continue;
                if(!d[e] && e != s) d[e] = d[u] + 1, q.push(e);
            }
        }
        return;
    };
    int d[N];
    memset(d, 0, sizeof(d));
    bfs(G, 0, d);
    
    // 4. 直径 の 両端 の 頂点 x, y は？
    // 4-1. 頂点 1 から 最も遠い 頂点 x を保存.
    int x = -1, r = 0;
    rep(i, N){
        if(r < d[i]){
            r = max(r, d[i]);
            x = i;
        }
    }
    assert(x != -1);
    
    // 4-2. 頂点 x から 各頂点までの距離を計算.
    int dx[N];
    memset(dx, 0, sizeof(dx));
    bfs(G, x, dx);
    
    // 4-3. 頂点 x から 最も遠い 頂点 y を保存.
    int y = -1;
    r = 0;
    rep(i, N){
        if(r < dx[i]){
            r = max(r, dx[i]);
            y = i;
        }
    }
    assert(y != -1);
    
    // 4-4. 頂点 y から 各頂点までの距離を計算.
    int dy[N];
    memset(dy, 0, sizeof(dy));
    bfs(G, y, dy);
    
    // 4-5. 直径 D は？
    int D = dx[y];
    
    // 5. 良さが, ある一定値以下となる塗り方をチェック.
    // 5-1. 良さとして存在しうる 最小値 を 保存.
    // -> 入力例 2 のように, 頂点 3 は, 直径 の 両端である 頂点 1, 5 からの距離が, 
    // ともに 2 であるが, 塗り方の良さが, 2 とはならないことに注意.
    // 理由として, 頂点 6 が, 直径 の 両端からの距離が, 3 なので, 
    // この場合, 頂点 3 は, 頂点 6 の影響を受ける形となり, 塗り方の良さは, 3 として集計される.
    // このような現象を考慮するため, minBest を 算出する必要が生じた.
    int minBest = 0; // 良さとして存在しうる最小値.
    rep(i, N){
        // 良さとして存在しうる最小値を更新.
        if(i == x || i == y) continue;
        int best = min(dx[i], dy[i]);
        minBest = max(minBest, best);
    }
    
    // 5-2. 頂点 x, y からの 各頂点 までの 距離 を 保存.
    // -> 良さが p以下となるには,  x, y からの距離が, ともに p以下
    vvi dxy(D + 1);
    rep(i, N){
        if(i == x || i == y) continue;
        int t = max({dx[i], dy[i], minBest}); // 良さとして存在しうる最小値を考慮.
        dxy[t].pb(i); // x, y からの距離が, ともに t以下 の 頂点 i を 追加.
    }
    
    // 5-3. 累積和.
    int cSum[D + 2];
    LL mPow2Sum[D + 2];
    memset(cSum, 0, sizeof(cSum));
    memset(mPow2Sum, 0, sizeof(mPow2Sum));
    rep(i, D + 1){
        cSum[i + 1] = cSum[i] + dxy[i].size();
        // 累積和 0 の 場合は, 2 の 冪乗 で カウントしないように集計する必要があることに注意.
        if(cSum[i + 1]) mPow2Sum[i + 1] = mPow2[cSum[i + 1]];
    }
    
    // 5-4. 特殊パターンを考慮(random_16.txtなど).
    // -> グラフが直線の場合, 何故か正しく動作しないので, 例外処理を追加.
    if(N > 2 && minBest * 2 + 2 == N) mPow2Sum[minBest + 1] = 1;
    
    // 6. 全ての塗り方に対する良さは？
    // 6-1. x, y が 同じ色.
    LL ans = 0;
    ans += 2LL * (LL)D * mPow2[N - 2]; // (x, y) = (白, 白), (黒, 黒)
    ans %= MOD;
    
    // 6-2. x, y が 異なる色.
    // -> 直径 の 両端 の 頂点 x, y 以外 の 頂点 z が 存在して, 
    // 頂点 x, z の 距離が, 直径 D に 等しくなるケースがありうるので, 
    // 上限 は 直径 D に 設定する必要があるはず.
    repx(i, minBest, D + 1){
        LL t = (mPow2Sum[i + 1] + MOD - mPow2Sum[i]) % MOD;
        LL u = (LL)i * t % MOD;
        u *= 2LL;
        u %= MOD;
        ans += u;
        ans %= MOD;
    }
    
    // 7. 出力.
    printf("%lld\n", ans);
    return 0;
    
}

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

// C++(GCC 9.2.1)

// 解き直し.

// https://atcoder.jp/contests/arc108/editorial/350

#include <bits/stdc++.h>

using namespace std;

using LL = long long;

using vi = vector<int>;

using vvi = vector<vi>;

using P = pair<int, int>;

#define repex(i, a, b, c) for(int i = a; i < b; i += c)

#define repx(i, a, b) repex(i, a, b, 1)

#define rep(i, n) repx(i, 0, n)

#define repr(i, a, b) for(int i = a; i >= b; i--)

#define pb push_back

const LL MOD = 1e9 + 7;

LL mPow2[202020]; // 2 の冪乗(MOD版)

int main(){

// 1. 入力情報.

int N;

scanf("%d", &N);

vvi G(N);

rep(i, N - 1){

int a, b;

scanf("%d %d", &a, &b);

a--;

b--;

G[a].pb(b);

G[b].pb(a);

}

// 2. 2 の冪乗(MOD版) を保存.

mPow2[0] = 1;

repx(i, 1, N) mPow2[i] = mPow2[i - 1] << 1, mPow2[i] %= MOD;

// 3. 頂点 1 から各頂点までの距離を計算.

// https://ja.wikipedia.org/wiki/幅優先探索

auto bfs = [&](vvi &G, int s, int* d){

// 3-1. 空のキュー.

queue<int> q;

// 3-2. 探索地点 s をキュー q に追加.

q.push(s);

while(!q.empty()){

// 3-3. キューから取り出す.

int u = q.front();

q.pop();

// 3-4. 取り出した要素を処理.

for(auto &e : G[u]){

// 3-5. 訪問済であれば, 処理をスキップ.

if(d[e]) continue;

if(!d[e] && e != s) d[e] = d[u] + 1, q.push(e);

}

return;

};

int d[N];

memset(d, 0, sizeof(d));

bfs(G, 0, d);

// 4. 直径の両端の頂点 x, y は？

// 4-1. 頂点 1 から最も遠い頂点 x を保存.

int x = -1, r = 0;

rep(i, N){

if(r < d[i]){

r = max(r, d[i]);

x = i;

}

assert(x != -1);

// 4-2. 頂点 x から各頂点までの距離を計算.

int dx[N];

memset(dx, 0, sizeof(dx));

bfs(G, x, dx);

// 4-3. 頂点 x から最も遠い頂点 y を保存.

int y = -1;

r = 0;

rep(i, N){

if(r < dx[i]){

r = max(r, dx[i]);

y = i;

}

assert(y != -1);

// 4-4. 頂点 y から各頂点までの距離を計算.

int dy[N];

memset(dy, 0, sizeof(dy));

bfs(G, y, dy);

// 4-5. 直径 D は？

int D = dx[y];

// 5. 良さが, ある一定値以下となる塗り方をチェック.

// 5-1. 良さとして存在しうる最小値を保存.

// -> 入力例 2 のように, 頂点 3 は, 直径の両端である頂点 1, 5 からの距離が,

// ともに 2 であるが, 塗り方の良さが, 2 とはならないことに注意.

// 理由として, 頂点 6 が, 直径の両端からの距離が, 3 なので,

// この場合, 頂点 3 は, 頂点 6 の影響を受ける形となり, 塗り方の良さは, 3 として集計される.

// このような現象を考慮するため, minBest を算出する必要が生じた.

int minBest = 0; // 良さとして存在しうる最小値.

rep(i, N){

// 良さとして存在しうる最小値を更新.

if(i == x || i == y) continue;

int best = min(dx[i], dy[i]);

minBest = max(minBest, best);

}

// 5-2. 頂点 x, y からの各頂点までの距離を保存.

// -> 良さが p以下となるには, x, y からの距離が, ともに p以下

vvi dxy(D + 1);

rep(i, N){

if(i == x || i == y) continue;

int t = max({dx[i], dy[i], minBest}); // 良さとして存在しうる最小値を考慮.

dxy[t].pb(i); // x, y からの距離が, ともに t以下の頂点 i を追加.

}

// 5-3. 累積和.

int cSum[D + 2];

LL mPow2Sum[D + 2];

memset(cSum, 0, sizeof(cSum));

memset(mPow2Sum, 0, sizeof(mPow2Sum));

rep(i, D + 1){

cSum[i + 1] = cSum[i] + dxy[i].size();

// 累積和 0 の場合は, 2 の冪乗でカウントしないように集計する必要があることに注意.

if(cSum[i + 1]) mPow2Sum[i + 1] = mPow2[cSum[i + 1]];

}

// 5-4. 特殊パターンを考慮(random_16.txtなど).

// -> グラフが直線の場合, 何故か正しく動作しないので, 例外処理を追加.

if(N > 2 && minBest * 2 + 2 == N) mPow2Sum[minBest + 1] = 1;

// 6. 全ての塗り方に対する良さは？

// 6-1. x, y が同じ色.

LL ans = 0;

ans += 2LL * (LL)D * mPow2[N - 2]; // (x, y) = (白, 白), (黒, 黒)

ans %= MOD;

// 6-2. x, y が異なる色.

// -> 直径の両端の頂点 x, y 以外の頂点 z が存在して,

// 頂点 x, z の距離が, 直径 D に等しくなるケースがありうるので,

// 上限は直径 D に設定する必要があるはず.

repx(i, minBest, D + 1){

LL t = (mPow2Sum[i + 1] + MOD - mPow2Sum[i]) % MOD;

LL u = (LL)i * t % MOD;

u *= 2LL;

u %= MOD;

ans += u;

ans %= MOD;

}

// 7. 出力.

printf("%lld\n", ans);

return 0;

}

[入力例]
2
1 2

[出力例]
2
※AtCoderのテストケースより

[入力例]
6
1 2
2 3
3 4
4 5
3 6

[出力例]
224
※AtCoderのテストケースより

[入力例]
35
25 4
33 7
11 26
32 4
12 7
31 27
19 6
10 22
17 12
28 24
28 1
24 15
30 24
24 11
23 18
14 15
4 29
33 24
15 34
11 3
4 35
5 34
34 2
16 19
7 18
19 31
22 8
13 26
20 6
20 9
4 33
4 8
29 19
15 21

[出力例]
298219707
※AtCoderのテストケースより

[入力例]
3
1 2
2 3

[出力例]
12

[入力例]
4
1 2
2 3
3 4

[出力例]
38

[入力例]
5
1 2
2 3
3 4
4 5

[出力例]
108

[入力例]
24
1 2
2 3
3 4
3 8
3 9
3 5
5 6
6 7
6 10
6 13
10 11
11 12
10 14
14 15
14 16
16 17
16 18
18 19
17 20
20 21
21 22
17 23
23 24

[出力例]
176021376

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

[入力例]

1 2

[出力例]

※AtCoderのテストケースより

[入力例]

1 2

2 3

3 4

4 5

3 6

[出力例]

224

※AtCoderのテストケースより

[入力例]

25 4

33 7

11 26

32 4

12 7

31 27

19 6

10 22

17 12

28 24

28 1

24 15

30 24

24 11

23 18

14 15

4 29

33 24

15 34

11 3

4 35

5 34

34 2

16 19

7 18

19 31

22 8

13 26

20 6

20 9

4 33

4 8

29 19

15 21

[出力例]

298219707

※AtCoderのテストケースより

[入力例]

1 2

2 3

[出力例]

[入力例]

1 2

2 3

3 4

[出力例]

[入力例]

1 2

2 3

3 4

4 5

[出力例]

108

[入力例]

1 2

2 3

3 4

3 8

3 9

3 5

5 6

6 7

6 10

6 13

10 11

11 12

10 14

14 15

14 16

16 17

16 18

18 19

17 20

20 21

21 22

17 23

23 24

[出力例]

176021376

■参照サイト
AtCoder Regular Contest 108

C++の練習を兼ねて, AtCoder Regular Contest 108 の 問題F (Paint Tree) を解いてみた.

コメントを残す コメントをキャンセル

C++の練習を兼ねて, AtCoder Regular Contest 108 の問題F (Paint Tree) を解いてみた.

コメントを残すコメントをキャンセル