結果

問題 No.121 傾向と対策:門松列(その2)
ユーザー lam6er
提出日時 2025-03-26 15:46:20
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 4,451 ms / 5,000 ms
コード長 3,135 bytes
コンパイル時間 296 ms
コンパイル使用メモリ 82,432 KB
実行使用メモリ 297,548 KB
最終ジャッジ日時 2025-03-26 15:46:46
合計ジャッジ時間 16,163 ms
ジャッジサーバーID
(参考情報)
judge3 / judge1
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
other AC * 9
権限があれば一括ダウンロードができます

ソースコード

diff #

class FenwickTree:
    def __init__(self, size):
        self.size = size
        self.tree = [0] * (self.size + 2)  # 1-based indexing

    def update(self, idx, delta):
        idx += 1  # converting 0-based to 1-based
        while idx <= self.size + 1:
            self.tree[idx] += delta
            idx += idx & -idx

    def query(self, idx):
        idx += 1  # converting 0-based to 1-based
        res = 0
        while idx > 0:
            res += self.tree[idx]
            idx -= idx & -idx
        return res

    def query_total(self):
        return self.query(self.size)

def main():
    import sys
    input = sys.stdin.read().split()
    n = int(input[0])
    A = list(map(int, input[1:n+1]))

    # Coordinate compression
    sorted_unique = sorted(set(A))
    compress_dict = {x: i for i, x in enumerate(sorted_unique)}
    compressed_A = [compress_dict[x] for x in A]
    m = len(sorted_unique)

    # Initialize count_right and right_fenwick
    count_right = [0] * m
    for x in compressed_A:
        count_right[x] += 1

    right_fenwick = FenwickTree(m)
    for x in compressed_A:
        right_fenwick.update(x, 1)

    # Initialize left_fenwick and count_left
    left_fenwick = FenwickTree(m)
    count_left = [0] * m

    # Initialize prod_fenwick with count_left[x] * count_right[x]
    prod_fenwick = FenwickTree(m)
    for x in range(m):
        initial_prod = count_left[x] * count_right[x]
        prod_fenwick.update(x, initial_prod)

    answer = 0

    for j in range(n):
        x = compressed_A[j]

        # Step 1: Remove x from right_fenwick and update count_right
        old_count_r = count_right[x]
        count_right[x] -= 1
        right_fenwick.update(x, -1)

        # Update prod_fenwick for x due to count_right change
        old_prod = count_left[x] * old_count_r
        new_prod = count_left[x] * count_right[x]
        delta_prod = new_prod - old_prod
        prod_fenwick.update(x, delta_prod)

        # Step 2: Compute L_less, L_greater, R_less, R_greater
        L_less = left_fenwick.query(x - 1)
        L_greater = left_fenwick.query_total() - left_fenwick.query(x)
        R_less = right_fenwick.query(x - 1)
        R_greater = right_fenwick.query_total() - right_fenwick.query(x)

        # Step 3: Compute same_pairs_less and same_pairs_greater
        same_pairs_less = prod_fenwick.query(x - 1)
        same_pairs_greater = prod_fenwick.query_total() - prod_fenwick.query(x)

        # Step 4: Calculate conditions and update answer
        cond1 = L_less * R_less - same_pairs_less
        cond2 = L_greater * R_greater - same_pairs_greater
        answer += max(0, cond1) + max(0, cond2)

        # Step 5: Add x to left_fenwick and update count_left
        old_count_l = count_left[x]
        count_left[x] += 1
        left_fenwick.update(x, 1)

        # Update prod_fenwick for x due to count_left change
        old_prod_l = old_count_l * count_right[x]
        new_prod_l = count_left[x] * count_right[x]
        delta_prod_l = new_prod_l - old_prod_l
        prod_fenwick.update(x, delta_prod_l)

    print(answer)

if __name__ == '__main__':
    main()
0