round_2015_qual_c_dijkstra.py

round_2015_qual_c_dijkstra.py

import functools
import sys

Q_1 = 1
Q_I = 'i'
Q_J = 'j'
Q_K = 'k'

QUATERNION_MUL = {
    (Q_1, Q_1): (+1, Q_1),
    (Q_1, Q_I): (+1, Q_I),
    (Q_1, Q_J): (+1, Q_J),
    (Q_1, Q_K): (+1, Q_K),

    (Q_I, Q_1): (+1, Q_I),
    (Q_I, Q_I): (-1, Q_1),
    (Q_I, Q_J): (+1, Q_K),
    (Q_I, Q_K): (-1, Q_J),

    (Q_J, Q_1): (+1, Q_J),
    (Q_J, Q_I): (-1, Q_K),
    (Q_J, Q_J): (-1, Q_1),
    (Q_J, Q_K): (+1, Q_I),

    (Q_K, Q_1): (+1, Q_K),
    (Q_K, Q_I): (+1, Q_J),
    (Q_K, Q_J): (-1, Q_I),
    (Q_K, Q_K): (-1, Q_1),
}


@functools.lru_cache(maxsize=None)
def full_reduce(sign, l):
    if len(l) == 1:
        return sign, l[0]
    if len(l) == 2:
        x = QUATERNION_MUL[(l[0], l[1])]
        return sign * x[0], x[1]
    left = full_reduce(1, l[:len(l)//2])
    right = full_reduce(1, l[len(l)//2:])
    res = QUATERNION_MUL[left[1], right[1]]
    return left[0] * right[0] * sign * res[0], res[1]


def can_make_ijk(s, n):
    # easy checks
    small = full_reduce(1, s)
    if small == (1, Q_1):
        return False
    if small[0] == 1 and n % 2 == 1:
        return False
    if small == (-1, Q_1) and n % 2 == 0:
        return False

    # truncate to relevant cycles
    n = min(n % 4 + 100, n)
    expanded = s * n
    if full_reduce(1, expanded) != (-1, Q_1):
        return False

    # do work
    right_excluded = set()
    for left in range(1, len(expanded) - 1):
        if full_reduce(1, tuple(expanded[:left])) == (1, Q_I):
            for right in range(len(expanded) - 1, left, -1):
                if right in right_excluded:
                    continue
                elif full_reduce(1, tuple(expanded[right:])) != (1, Q_K):
                    right_excluded.add(right)
                elif full_reduce(1, tuple(expanded[left:right])) == (1, Q_J):
                    return True
    return False


if __name__ == '__main__':
    specific = int(sys.argv[2]) if len(sys.argv) > 2 else None
    with open(sys.argv[1], 'r') as f:
        amount = int(f.readline())
        for i in range(1, amount + 1):
            _, n = f.readline().split()
            s = f.readline().strip()
            if specific is None or i == specific:
                word = {True: 'YES', False: 'NO'}[can_make_ijk(s, int(n))]
                print(f'CASE #{i}: {word}')