ABC014 B - 価格の合計

備忘録

問題

atcoder.jp

回答

import sys, os, math, bisect, itertools, collections, heapq, queue
# from scipy.sparse.csgraph import csgraph_from_dense, floyd_warshall
from decimal import Decimal
from collections import defaultdict, deque

sys.setrecursionlimit(10000000)

ii = lambda: int(sys.stdin.buffer.readline().rstrip())
il = lambda: list(map(int, sys.stdin.buffer.readline().split()))
fl = lambda: list(map(float, sys.stdin.buffer.readline().split()))
iln = lambda n: [int(sys.stdin.buffer.readline().rstrip()) for _ in range(n)]

iss = lambda: sys.stdin.buffer.readline().decode().rstrip()
sl = lambda: list(map(str, sys.stdin.buffer.readline().decode().split()))
isn = lambda n: [sys.stdin.buffer.readline().decode().rstrip() for _ in range(n)]

lcm = lambda x, y: (x * y) // math.gcd(x, y)

MOD = 10 ** 9 + 7
MAX = float('inf')


def main():
    if os.getenv("LOCAL"):
        sys.stdin = open("input.txt", "r")

    N, X = il()
    A = il()

    ret = 0
    for n in range(N):
        if X >> n & 1:
            ret += A[n]
    print(ret)


if __name__ == '__main__':
    main()

考え方

Xに対してnbit右シフトし、
その商品が部分集合に含まれるか否かを判定する。


Xを2進数表記した際に、bitが1になっている位置の商品が、
部分集合として含まれている。
つまり、X = 5のとき、2進数表記で0101なので、
0-indexedとすると、0番目と2番目の商品が部分集合に含まれているとする。

n番目の商品が部分集合に含まれているか否かは、
Xを右にnbitシフトした時、0桁目が1か否かで判定できる。
そのため、Xnビット右シフトした値と1andを取ることで判定する。

# X = 5, N = 4のとき
X >> 0 & 1 # 0b101 & 1 = True
X >> 1 & 1 # 0b10 & 1 = False
X >> 2 & 1 # 0b1 & 1 = True
X >> 3 & 1 # 0b0 & 1 = False

bit全探索っぽい