Python
Used sliding window for part 3. The off-by-one difference between i and j tripped me up for a while.
def part1(data: str, mentor: str = 'A'):
mentors = 0
pairs = 0
for p in data:
if p == mentor:
mentors += 1
elif p == mentor.lower():
pairs += mentors
return pairs
assert part1("ABabACacBCbca") == 5
def part2(data: str):
all_pairs = 0
for mentor in 'ABC':
all_pairs += part1(data, mentor)
return all_pairs
assert part2("ABabACacBCbca") == 11
from collections import defaultdict
def part3(data: str, distance_limit: int = 1000, repeat: int = 1000):
n = len(data)
N = len(data) * repeat
pairs = 0
person_count = defaultdict(int)
curr = 0
# initialize the first window excluding the right boundary
for j in range(distance_limit):
person_count[data[j % n]] += 1
for curr in range(N):
# move left boundary (if applicable)
if (i := curr - distance_limit - 1) >= 0:
person_count[data[i % n]] -= 1
# move right boundary (if applicable)
if (j := curr + distance_limit) < N:
person_count[data[j % n]] += 1
# if mentee, record pairs
if data[curr % n].islower():
pairs += person_count[data[curr % n].upper()]
return pairs
assert (t := part3("AABCBABCABCabcabcABCCBAACBCa", 10, 1)) == 34, f"Expected 34 but got {t}"
assert (t := part3("AABCBABCABCabcabcABCCBAACBCa", 10, 2)) == 72, f"Expected 72 but got {t}"
assert (t := part3("AABCBABCABCabcabcABCCBAACBCa", 1000, 1000)) == 3442321, f"Expected 3442321 but got {t}"-
Python