2021-07-07

AtCoder ABC105: D - Candy Distribution

競技プログラミング AtCoder

区間和は「累積和の差分」で考えてみましょう。水色diff です。

問題

数列 $A_1, \cdots, A_{N}$ の区間 $(l, r)$ で $A_l + A_{l+1} + \cdots + A_r$ が $M$ の倍数であるものの数を求めよ。

制約

$1 \le N \le 10^5$
$1 \le M \le 10^5$
$1 \le A_i \le 10^9$

考察

$\{A_i\}$ の累積和 $\{S_i\}$ を計算し、

$\begin{aligned} A_l + A_{l+1} + \cdots + A_r = \text{(}M\text{の倍数)} &\Longleftrightarrow S_{r} - S_{l-1} = \text{(}M\text{の倍数)} \\ &\Longleftrightarrow S_{r} \equiv S_{l-1} \quad \text{(mod} M \text{)} \end{aligned}$

とします。

求める値は、 $\{S_i\}$ の剰余の数列 $\{T_i\}$ で、同じ値が何度出てくるかの総和と等しくなります。

Pythonでは collections.Counter が便利です。 values() で出現回数が取得できます。

実装例

N, M = map(int, input().split())
A = list(map(int, input().split()))

S = [0] + list(itertools.accumulate(A))
T = [s % M for s in S]
counter = collections.Counter(T)
ans = sum(n * (n-1) // 2 for n in counter.values())
print(ans)

2021-07-06

AtCoder 第六回アルゴリズム実技検定 (PAST): N - 活動

競技プログラミング AtCoder

「条件を手で計算してみる大切さ」を教えてくれる問題です。ソートしてからナップサックDPで解けます。

問題

$N$ 種類の活動の中から $1$ つ以上を選び、好きな順序で行なう。最初の体力は $H$ であり、活動 $i$ を行うと、得点 $a_i \times \text{(体力)}$ を得た後、体力が $b_i$ 減少する。得られる得点の和の最大値を求めよ。

制約

$1 \le N \le 100$
$1 \le H \le 10^5$
$1 \le a_i, b_i \le 10^5$

考察

活動を行った時に得られる得点が「その時点の体力」に依存するので、そのままナップサックDPは使えません。しかし、活動の順序を全検索（ $100!$ 通り）することは、時間的に不可能です。

ここで、活動 $1$ → 活動 $2$ を行った場合と、活動 $1$ → 活動 $2$ を行った場合の、それぞれ得られる得点を計算してみます。体力の減少は、どちらも同じ $b_1+b_2$ です。

$\displaystyle \begin{aligned} & \left\{ \begin{array}{l} 体力 \, h \\ 得点 \, 0 \end{array} \right. & \xrightarrow{活動1} \, & \left\{ \begin{array}{l} 体力 \, h-b_1 \\ 得点 \, a_1 h \end{array} \right. & \xrightarrow{活動2} \, & \left\{ \begin{array}{l} 体力 \, h-b_1-b_2 \\ 得点 \, a_1 h + a_2(h-b_1) \end{array} \\ \right. \\ & \left\{ \begin{array}{l} 体力 \, h \\ 得点 \, 0 \end{array} \right. & \xrightarrow{活動2} \, & \left\{ \begin{array}{l} 体力 \, h-b_2 \\ 得点 \, a_2 h \end{array} \right. & \xrightarrow{活動1} \, & \left\{ \begin{array}{l} 体力 \, h-b_2-b_1 \\ 得点 \, a_2 h + a_1(h-b_2) \end{array} \\ \right. \end{aligned}$

得られる得点の差は、

$\displaystyle \begin{aligned} \text{score}_{\text{活動1} \rightarrow \text{活動2}} - \text{score}_{\text{活動2} \rightarrow \text{活動1}} &= \left\{a_1 h + a_2(h-b_1) \right\} - \left\{a_2 h + a_1(h-b_2) \right\} \\ &= a_1 b_2 - a_2 b_1 \end{aligned}$

それで、

$\displaystyle \begin{aligned} \text{score}_{\text{活動1} \rightarrow \text{活動2}} \ge \text{score}_{\text{活動2} \rightarrow \text{活動1}} & \Longleftrightarrow a_1 b_2 - a_2 b_1 \ge 0 \\ & \Longleftrightarrow \frac{a_1}{b_1} \ge \frac{a_2}{b_2} \\ \end{aligned}$

となります。ゆえに、 $\displaystyle \frac{a_i}{b_i}$ が大きい方から処理する（降順でソート）のが最善です。

実装例

INF = 10**10
 
N, H = map(int, input().split())
 
Activities = []
for i in range(N):
    a, b = map(int, input().split())
    Activities.append((a, b))
Activities.sort(key=lambda x: (x[0]/x[1]), reverse=True)
 
dp = [[-INF] * (H+1) for _ in range(N+1)]
dp[0][H] = 0

for i in range(N):
    A, B = Activities[i]
    for h in range(H+1):
        dp[i+1][h] = max(dp[i+1][h], dp[i][h]) 
        if h > 0:
            h1 = max(h-B, 0)
            dp[i+1][h1] = max(dp[i+1][h1], dp[i][h] + A*h) 

print(max(dp[N]))

DP高速化

ここからは趣味の範囲です。

このコードで、PyPy で 336 ms でACできますが、さらに高速化できます。

$\text{dp}_{i+1}$ は $\text{dp}_i$ にしか依存しないので、 $\text{dp}_0,\cdots,\text{dp}_N$ を管理するのではなく、 $\text{dp}_i$ （dpとする）と $\text{dp}_{i+1}$ （dp1とする）の2つのみ管理するようにします。ループごとに dp1 を dp で更新し、ループの最後に dp を dp1 に置き換えます。

dp = [-INF] * (H+1)
dp[H] = 0

for A, B in Activities:
    dp1 = [-INF] * (H+1)
    for h in range(H+1):
        dp1[h] = max(dp1[h], dp[h]) 
        if h > 0:
            h1 = max(h-B, 0)
            dp1[h1] = max(dp1[h1], dp[h] + A*h) 
    dp = dp1

print(max(dp))

これで、PyPy で 260 ms です。

さらに、配るDPの更新の時に大きい $h$ のみを更新しているので、 $\text{dp}$ 1本を $h=0,1,\cdots,H$ で更新していけばOKです。

dp = [-INF] * (H+1)
dp[H] = 0

for A, B in Activities:
    for h in range(H+1):
        if h > 0:
            h1 = max(h-B, 0)
            dp[h1] = max(dp[h1], dp[h] + A*h) 

print(max(dp))

これで、PyPy で 150 ms です。

ただし、これでもPythonではTLEします。計算量 $O(HN)$ が $10^7$ なので、PyPyでないと時間が厳しいです。

2021-07-05

AtCoder ABC208 D - Shortest Path Queries 2

競技プログラミング AtCoder

ワーシャル・フロイド法の理解を問う問題です。緑diffです。

問題

$G$ を $N$ 頂点 $M$ 辺の有向グラフ、 $f(s,t,k)$ を $s$ から $t$ へ頂点番号 $k$ 以下の頂点のみ経由したときの最小距離とする。 $\displaystyle \sum_{s=1}^N \sum_{t=1}^N \sum_{k=1}^N f(s,t,k)$ を求めよ。

制約

$1 \le N \le 400$
$M \le N(N-1)$

考察

全点間の距離なので、計算量 $O(N ^3)$ のワーシャル・フロイド法が思いつきます。しかし、 $k=1,\cdots,N$ に対してワーシャル・フロイド法を用いると、全体で $O(N ^4)$ になり、間に合いません。

ここで、ワーシャル・フロイド法のコード

for k in range(N):
    for i in range(N):
        for j in range(N):
            if cost[i][k] != INF and cost[k][j] != INF:
                cost[i][j] = min(cost[i][j], cost[i][k] + cost[k][j])

を眺めると、一番外側の $k$ のループがまさに、頂点 $k$ を付加して全点間の距離を更新する処理になっていることが分かります。

それで $\text{cost}_{i, j}$ を、

$\text{cost}_{u, u} = 0 (u=1,\cdots,N)$
$\text{cost}_{A_i, B_i} = C_i$
$\text{cost}_{\text{other}, \text{other}} = \infty$

で初期化し、 $k$ の順に処理します。 $k$ の時の全点間の距離の和は差分処理で求めています。

実装例

INF = 10**15

N, M = map(int, input().split())
cost = [[INF] * N for _ in range(N)]
for i in range(N):
    cost[i][i] = 0

for _ in range(M):
    a, b, c = map(int, input().split())
    cost[a-1][b-1] = c

v = 0
for a in range(N):
    for b in range(N):
        v += (cost[a][b] if cost[a][b] < INF else 0)

ans = 0
for k in range(N):
    for a in range(N):
        for b in range(N):
            d = cost[a][k] + cost[k][b]
            if d < cost[a][b]:
                v += -(cost[a][b] if cost[a][b] < INF else 0) + d
                cost[a][b] = d

    ans += v

print(ans)