第12回 JOI 予選 E 魚の生息範囲【Python】

問題はこちら。

３次元の座標圧縮 & 累積和。コードは長いけどやってることは単純。

# 入力
N, K = map(int, input().split())
X, Y, D = [], [], []
plot = []
for _ in range(N):
   x_1, y_1, d_1, x_2, y_2, d_2 = map(int, input().split())
   X.append(x_1); X.append(x_2)
   Y.append(y_1); Y.append(y_2)
   D.append(d_1); D.append(d_2)
   plot.append((x_1, y_1, d_1, 1))
   plot.append((x_2, y_1, d_1, -1))
   plot.append((x_1, y_2, d_1, -1))
   plot.append((x_1, y_1, d_2, -1))
   plot.append((x_1, y_2, d_2, 1))
   plot.append((x_2, y_1, d_2, 1))
   plot.append((x_2, y_2, d_1, 1))
   plot.append((x_2, y_2, d_2, -1))


# 座標からindexを対応させる辞書を作る
X = list(set(X)); Y = list(set(Y)); D = list(set(D))
X.sort(); Y.sort(); D.sort()
indX = {x:i + 1 for i, x in enumerate(X)}
indY = {y:i + 1 for i, y in enumerate(Y)}
indD = {d:i + 1 for i, d in enumerate(D)}


# 累積和の初期化
S = [[[0]*(2*N + 1) for _ in range(2*N + 1)] for _ in range(2*N + 1)]
for x, y, d, w in plot:
   i, j, k = indX[x], indY[y], indD[d]
   S[i][j][k] += w
   

# 累積和の計算
for i in range(1, 2*N + 1):
   for j in range(1, 2*N + 1):
       for k in range(1, 2*N + 1):
           S[i][j][k] += S[i - 1][j][k] + S[i][j - 1][k] + S[i][j][k - 1]\
               - S[i][j - 1][k - 1] - S[i - 1][j][k - 1] - S[i - 1][j - 1][k]\
                   + S[i - 1][j - 1][k - 1]
                   

# 答えの計算       
ans = 0
for i in range(1, 2*N + 1):
   for j in range(1, 2*N + 1):
       for k in range(1, 2*N + 1):
           if S[i][j][k] >= K:
               ans += (X[i] - X[i - 1]) * (Y[j] - Y[j - 1]) * (D[k] - D[k - 1])


# 出力
print(ans)