D - Distinct Trio

[hxrxchang]

https://atcoder.jp/contests/abc252/tasks/abc252_d

問題概要

[]int のAから異なる要素を3つ取り出す組み合わせはいくつあるか？

解法

ナイーブな実装は

for i in range(N - 2):
  for j in range(i + 1, N - 1):
     for k in range(j + 1, N):
        if A[i] !=  A[j] and A[j] != A[k] and A[k] != A[i]:
            cnt += 1

で3重ループだが、Nは10**5なので間に合わないので計算量を減らす工夫が必要。

計算量削減

Aの中から要素を3つ見るというのは変わらず。 A[i] を考える。Aの中で、X < A[i] < Y となる、XとYの組み合わせがAの中でいくつあるかを求めれば、O(N)で全組み合わせを求められる。 Aの中でA[i]より小さい要素の数、A[i]より大きい要素の数は、

1. 0からmax(A)の各数字がAで何回出現するか(これをcntとする)
2. ↑の累積和

をとれば、cnt[A[i] - 1] * (N - cnt[A[i]]) で求められる。

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

# A[i] < A[j] < A[k] となるような組み合わせの数を求める
# A[j] について考えて、A[j] より小さい数と A[j] より大きい数の数を数えればOK
limit = max(A)
cnt = [0] * (limit + 1)
# 各要素の出現回数を求める
for a in A:
    cnt[a] += 1

# cnt[i]: Aにi以下の数がいくつあるかを計算
for i in range(limit):
    cnt[i + 1] += cnt[i]

ans = 0
for a in A:
    countLessThanA = cnt[a-1]
    countGreaterThanA = N - cnt[a]
    ans += countLessThanA * countGreaterThanA

print(ans)

提出: https://atcoder.jp/contests/abc252/submissions/51927592