개발자 톡

연습문제 톡 CPTI

파이썬 - 비트연산으로..

등록일
2025-02-28 18:01:50
조회수
57
작성자
Avantgarde95

스트링으로 했더니 타임아웃 떠서 정수로 바꾸고 비트연산으로 했더니 겨우 통과했네요..


import sys

def countEqualPairs(values):
    countMap = {}

    for value in values:
        countMap[value] = countMap.get(value, 0) + 1

    result = 0

    for count in countMap.values():
        result += (count * (count - 1)) // 2

    return result

def countEqualPairsFrom2(values1, values2):
    countMap = {}

    for value in values1:
        if value in countMap:
            countMap[value][0] += 1
        else:
            countMap[value] = [1, 0]

    for value in values2:
        if value in countMap:
            countMap[value][1] += 1
        else:
            countMap[value] = [0, 1]

    result = 0
    
    for a, b in countMap.values():
        result += a * b

    return result

# ===================

N, M = map(int, sys.stdin.readline().rstrip().split())
inputs = []

for i in range(N):
    string = sys.stdin.readline().rstrip()
    value = int(string, 2)
    inputs.append(value)

# ===================

pow2 = [1]

for i in range(1, M + 1):
    pow2.append(pow2[i - 1] * 2)

def binaryDigit(value, i):
    return (value & pow2[i]) >> i;

# ===================

answer = 0

# 모두 같은 경우.
answer += countEqualPairs(inputs)

# 하나만 다른 경우.
for i in range(M):
    listMap = {0: [], 1: []}
    compareMask = pow2[M] - 1 - pow2[i] # i번째만 빼고 비교.

    for input in inputs:
        key = binaryDigit(input, i)
        listMap[key].append(input & compareMask)

    answer += countEqualPairsFrom2(listMap[0], listMap[1])

# 두개가 다른 경우.
for i1 in range(M):
    for i2 in range(i1 + 1, M):
        listMap = {(0, 0): [], (0, 1): [], (1, 0): [], (1, 1): []}
        compareMask = pow2[M] - 1 - pow2[i1] - pow2[i2] # i1, i2번째만 빼고 비교.

        for input in inputs:
            key = (binaryDigit(input, i1), binaryDigit(input, i2))
            listMap[key].append(input & compareMask)
        
        answer += countEqualPairsFrom2(listMap[(0, 0)], listMap[(1, 1)])
        answer += countEqualPairsFrom2(listMap[(0, 1)], listMap[(1, 0)])

print(answer)
#CPTI

이 카테고리의 톡 더보기