dna.cpp

#include <iostream>
#include <vector>

struct Count {
  int at, ca, tc;
  int qw;

  Count() : at(0), ca(0), tc(0), qw(0) {}

  inline Count operator-(Count rhs) const {
    Count result = *this;
    result.at -= rhs.at;
    result.ca -= rhs.ca;
    result.tc -= rhs.tc;
    result.qw -= rhs.qw;
    return result;
  }
};

std::vector<Count> prefixes;

void init(std::string a, std::string b) {
  prefixes.reserve(a.size() + 1);
  Count count;
  prefixes.push_back(count);
  for (int i = 0; i < a.size(); i++) {
    if (a[i] == 'A' && b[i] == 'T') count.at += 1;
    if (a[i] == 'T' && b[i] == 'A') count.at -= 1;
    if (a[i] == 'C' && b[i] == 'A') count.ca += 1;
    if (a[i] == 'A' && b[i] == 'C') count.ca -= 1;
    if (a[i] == 'T' && b[i] == 'C') count.tc += 1;
    if (a[i] == 'C' && b[i] == 'T') count.tc -= 1;
    if (a[i] != b[i]) count.qw += 1;
    prefixes.push_back(count);
  }
}

int get_distance(int x, int y) {
  y += 1;

  Count count = prefixes[y] - prefixes[x];

  if (!(count.at == count.ca && count.ca == count.tc)) return -1;

  return (count.qw + std::abs(count.at)) / 2;
}

dna.py

import copy
from dataclasses import dataclass


@dataclass
class Count:
    at: int
    ca: int
    tc: int
    qw: int

    def __sub__(self, rhs: "Count") -> "Count":
        return Count(
            at = self.at - rhs.at,
            ca = self.ca - rhs.ca,
            tc = self.tc - rhs.tc,
            qw = self.qw - rhs.qw,
        )


def init(a: str, b: str) -> None:
    global prefixes

    count = Count(0, 0, 0, 0)
    prefixes = [copy.copy(count)]
    for i in range(len(a)):
        if a[i] == 'A' and b[i] == 'T': count.at += 1
        if a[i] == 'T' and b[i] == 'A': count.at -= 1
        if a[i] == 'C' and b[i] == 'A': count.ca += 1
        if a[i] == 'A' and b[i] == 'C': count.ca -= 1
        if a[i] == 'T' and b[i] == 'C': count.tc += 1
        if a[i] == 'C' and b[i] == 'T': count.tc -= 1
        if a[i] != b[i]: count.qw += 1
        prefixes.append(copy.copy(count))


def get_distance(x: int, y: int) -> int:
    y += 1

    count = prefixes[y] - prefixes[x]

    if not count.at == count.ca == count.tc: return -1

    return (count.qw + abs(count.at)) // 2