CODE FESTIVAL 2016 Final B - Exactly N points

備忘録

問題

atcoder.jp

回答

import sys
import os
import math
import bisect
import itertools
import collections
import heapq
import queue
import array
import time
import numpy as np
from numpy import linalg as LA

# 時々使う
# import numpy as np
# from decimal import Decimal, ROUND_HALF_UP
# from scipy.sparse.csgraph import csgraph_from_dense, floyd_warshall
# from collections import defaultdict, deque

# 再帰の制限設定
sys.setrecursionlimit(10000000)


def ii(): return int(sys.stdin.buffer.readline().rstrip())
def il(): return list(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 = ii()

    ret = set()
    for n in range(1, N+1):
        q = ((n + 1) * n) // 2
        if q == N:
            ret.add(n)
            break
        if q < N:
            ret.add(n)
        else:
            ret.discard(q-N)
            ret.add(n)
            break

    print(*list(ret), sep="\n")


if __name__ == '__main__':
    main()

考え方

i番目の問題と解いたとき、得点の合計がNと同じか上回るとき、
ちょうどN点を取ることができます。

例えば、
1問目までの全ての問題を回答した時の合計点数は1です。
このとき、1点のみ取ることが出来ます。
2問目までの全ての問題を回答した時の合計点数は3です。
このとき、丁度1, 2, 3点を取ることができます。
3問目までの全ての問題を回答した時の合計得点は6です。
このとき、丁度1, 2, 3, 4, 5, 6点を取ることができます。

以上から、合計得点がNを上回るか同じとなる、iまでを解いた時点から、
不要な問題のみを省いた集合が回答となります。
具体的には、n(1 <= n <= N)問まで全て解いたときの合計点数(q)を和の公式で求め、
合計点数がちょうどN(q == N)の場合にはn問まで全て解くことが回答となり、
合計点数がNを上回っている場合(q > N)にはN - q点の問題を解く必要がありません。