ABC044 C - 高橋君とカード
備忘録
問題
回答
import sys import os import math import bisect import itertools import collections import heapq import queue import array import time import datetime # 時々使う # import numpy as np # from decimal import Decimal, ROUND_HALF_UP # from scipy.sparse.csgraph import csgraph_from_dense, floyd_warshall # 再帰の制限設定 sys.setrecursionlimit(10000000) def ii(): return int(sys.stdin.buffer.readline().rstrip()) def il(): return list(map(int, sys.stdin.buffer.readline().split())) def it(): return tuple(map(int, sys.stdin.buffer.readline().split())) def fl(): return list(map(float, sys.stdin.buffer.readline().split())) def iln(n): return [int(sys.stdin.buffer.readline().rstrip()) for _ in range(n)] def iss(): return sys.stdin.buffer.readline().decode().rstrip() def sl(): return list(map(str, sys.stdin.buffer.readline().decode().split())) def isn(n): return [sys.stdin.buffer.readline().decode().rstrip() for _ in range(n)] def lcm(x, y): return (x * y) // math.gcd(x, y) # MOD = 10 ** 9 + 7 MOD = 998244353 INF = float('inf') def main(): if os.getenv("LOCAL"): sys.stdin = open("input.txt", "r") N, A = il() X = il() sm = sum(X) + 1 dp = [[[0] * 2600 for _ in range(N + 1)] for _ in range(N + 1)] dp[0][0][0] = 1 for i in range(N): for j in range(N): for s in range(sm): # カードを選択する場合 # => i枚の中からj+1枚選んだとき、合計はs+X[i] dp[i+1][j+1][s+X[i]] += dp[i][j][s] # カードを選択しない場合 # => i毎の中からj枚を選んだまま。合計も変化なし。 dp[i+1][j][s] += dp[i][j][s] ret = 0 for j in range(1, N+1): ret += dp[N][j][A*j] print(ret) if __name__ == '__main__': main()
考え方
雰囲気だけ理解した問題。
ちなみに、上記回答はPyPy3 (7.3.0)で提出時のみ通る。
素朴に全てのカードに対して、使うか否かの全探索を行なうと、
最大でO(2 ** 50)となり、時間が足りない。
したがって、カードを選択した時の合計値の遷移を動的計画法で求める。
dp[i][j][s]をi枚のカードの中からj枚選択して、合計がsとなる選択数として、状態の遷移をメモする。
また遷移時にはカードを選択した場合と、しない場合の遷移を行う。
あとは求めた全ての状態から、
A * j(j枚選んだときの合計がA * j ＝ 平均がA）となる選択数を合算する。
参考
AtCoder ABC 044 C - 高橋君とカード (ARC 060 C) (水色, 300 点) - けんちょんの競プロ精進記録
