RSS preview of John D. Cook

Rss preview of Blog of John D. Cook

Forensic accounting in Python

2026-07-21 23:05:04

I recently had a project in which I had to reverse engineer a data analysis. There was some ambiguity regarding which of several possibilities someone chose for several of the variables, something analogous to the following example.

Suppose you have three numbers with uncertain values with a known, or at least purported, sum. The first number could be 31, 41, or 59; the second could be either 26 or 53; the last could be 58, 97, 93, or 23.

The following code enumerates all 3 × 2 × 4 = 24 possibilities and prints their sums.

from itertools import product

# Example input
possibilities = [(31, 41, 59), (26, 53), (58, 97, 93, 23)]

for combo in product(*possibilities):
    total = sum(combo) 
    print(f"Combination {combo} sums to: {total}")

In this example all the sums are unique, though of course that might not happen in practice. If, for example, you know the sum is 187, you know the three numbers were 41, 53, and 93. If the reported sum is 200, you know some assumption has been violated because none of the possible choices add up to 200.

Solutions

There are eight solutions, each being a flip or a rotation of one fundamental solution. Here’s a representative.

SWI Prolog code

:- use_module(library(clpfd)).

%% Enumerate all ways to place 5 white queens and 3 black queens
%% on a 5x5 board so that no queen attacks a queen of the opposite colour.
%% Same-colour queens are allowed to attack each other.
%%
%% Positions are numbered 1..25 in row-major order:
%%   1  2  3  4  5
%%   6  7  8  9 10
%%  11 12 13 14 15
%%  16 17 18 19 20
%%  21 22 23 24 25

solve :-
    solution(Whites, Blacks),
    print_solution(Whites, Blacks),
    fail.
solve :-
    writeln('No more solutions.').

solution(Whites, Blacks) :-
    length(Whites, 5),
    length(Blacks, 3),
    append(Whites, Blacks, All),
    All ins 1..25,
    all_distinct(All),
    % canonical order so each set of positions is generated once
    Whites = [W1,W2,W3,W4,W5],
    W1 #< W2, W2 #< W3, W3 #< W4, W4 #< W5,
    Blacks = [B1,B2,B3],
    B1 #< B2, B2 #< B3, % every black is safe from every white maplist(safe_from_all(Whites), Blacks), label(All). safe_from_all(Whites, Black) :- maplist(no_attack(Black), Whites). %% Two queens do *not* attack each other. no_attack(P, Q) :- R1 #= (P-1) div 5 + 1, C1 #= (P-1) mod 5 + 1, R2 #= (Q-1) div 5 + 1, C2 #= (Q-1) mod 5 + 1, R1 #\= R2, % different row C1 #\= C2, % different column abs(R1-R2) #\= abs(C1-C2). % different diagonal print_solution(Whites, Blacks) :- format('Whites: ~w~n', [Whites]), format('Blacks: ~w~n', [Blacks]), forall(between(1, 5, R), ( forall(between(1, 5, C), ( Pos is (R-1)*5 + C, ( memberchk(Pos, Whites) -> write('W ')
                      ; memberchk(Pos, Blacks) -> write('B ')
                      ; write('. ')
                      )
                    )),
             nl )),
    nl.

Lean 4 code

-- Lean 4 solution for the 5×5 black-and-white queens puzzle.
-- 5 white queens + 3 black queens; no opposite-colour attacks.

def absDiff (a b : Nat) : Nat :=
  if a ≥ b then a - b else b - a

def attacks (p q : Nat × Nat) : Bool :=
  p.1 == q.1 || p.2 == q.2 || absDiff p.1 q.1 == absDiff p.2 q.2

/-- All 25 squares (0-based, row-major). -/
def allPositions : List (Nat × Nat) :=
  (List.range 5).flatMap fun r =>
    (List.range 5).map fun c => (r, c)

/-- Combinations of length exactly `k` (order-preserving). -/
def combinations {α : Type} : Nat → List α → List (List α)
  | 0,     _      => [[]]
  | _,     []     => []
  | k + 1, x::ys  =>
      (combinations k ys).map (fun zs => x :: zs) ++
      combinations (k + 1) ys

def isAttackedBy (pos : Nat × Nat) (queens : List (Nat × Nat)) : Bool :=
  queens.any (fun q => attacks q pos)

/-- Every solution as a pair (white squares, black squares). -/
def solve : List (List (Nat × Nat) × List (Nat × Nat)) :=
  (combinations 3 allPositions).foldl (fun acc blacks =>
    let forbidden := allPositions.filter fun p =>
      blacks.contains p || isAttackedBy p blacks
    let available := allPositions.filter fun p =>
      !(forbidden.contains p)
    acc ++ (combinations 5 available).map fun whites => (whites, blacks)
  ) []

/-- Pretty-print one board. -/
def printBoard (whites blacks : List (Nat × Nat)) : IO Unit := do
  for r in List.range 5 do
    let mut row := ""
    for c in List.range 5 do
      let p := (r, c)
      row := row ++
        (if whites.contains p then "W "
         else if blacks.contains p then "B "
         else ". ")
    IO.println row
  IO.println ""

/-- Enumerate and display every solution. -/
def main : IO Unit := do
  let sols := solve
  IO.println s!"Found {sols.length} solutions.\n"
  for p in (List.range sols.length).zip sols do
    let i := p.1
    let (whites, blacks) := p.2
    IO.println s!"Solution {i + 1}:"
    printBoard whites blacks

#eval main

The post Solving a chess puzzle with Grok 4.5 first appeared on John D. Cook.

Fitting a regular expression to a list of words

2026-07-20 03:30:17

Suppose you want to search for a list of words. If you’re using grep, you can add the -f flag provide a file of regular expressions, and you can add the -F to tell it that the regular expressions are in fact just words. I did something like this a couple days ago when searching for diagnosis codes.

grep -w -F -o -f icd10codes.txt notes.txt

Now you might want to combine your list of words into a singular regular expression, for efficiency or possibly for some other reason. Apparently ripgrep does this because when I tried replacing grep with ripgrep in the command above I got an error saying “Compiled regex exceeds size limit of 104857600 bytes.”

Beating brute force

Say you wanted to search for the strings “bluecross”, “blueshield”, and “bluey”. You could simply form the brute force regular expression

bluecross|blueshied|bluey

but that doesn’t take advantage of the fact that all three strings begin with “blue.” A smaller regular expression would be

blue(shield|cross|y)

Finding the shortest regular expression that matches a list of words is a hard problem, but finding a regular expression that’s shorter than brute force is not. The Python package trieregex will do this. According to the documentation,

trieregex creates efficient regular expressions (regexes) by storing a list of words in a trie structure, and translating the trie into a more compact pattern.

Let’s try our blue example with trieregex.

import re
from trieregex import TrieRegEx as TRE

words = ['bluecross', 'blueshield', 'bluey']
tre = TRE(*words) 
print(tre.regex())

This produces the same regular expression as above, except it adds ?: to make the parentheses non-capturing.

blue(?:shield|cross|y)

Prefixes versus suffixes

The library builds a trie data structure using common prefixes. That works well in the example above, but the result is disappointing when we have common suffixes rather than common prefixes. The following code

words = ['javascript', 'typescript']
tre = TRE(*words) 
print(tre.regex())

produces the regular expression

(?:javascript|typescript)

which is no better than brute force, whereas we might have hoped for

(?:java|type)script

HCPCS examples

As mentioned at the top of the post, ripgrep failed to search on a list of ICD-10 codes. The list of HCPCS codes is about 10x smaller, and more compressible. Ripgrep was able to fit all HCPCS codes into a single regex and was able to search the test file much faster than grep. The command

grep -w -F -o -f hcpcs.txt notes.txt

took 73.426 seconds to execute, while the command

rg -w -F -o -f hcpsc.txt notes.txt

took 0.078 seconds, three orders of magnitude faster.

The following code will read a list of HCPCS codes from a file and create a regular expression.

tre = TRE()
with open('hcpcs.txt', 'r') as file:
    for line in file:
        tre.add(line.strip())
print(len(tre.regex()))

This shows that the resulting regular expression has 17,198 characters. The file of codes has 8725 five-character codes, so the regex compresses the code characters by roughly a ratio of 5 to 2.

The post Fitting a regular expression to a list of words first appeared on John D. Cook.

Sum of low squares

2026-07-20 00:41:48

Squares, high and low

Let p be an odd prime number. Then half the numbers from 1 through p − 1 are squares and half are not. That is, for half of numbers 1 ≤ k < p, the equation

x² = k mod p

has a solution. The traditional name for these numbers is “quadratic residues” but we can just say “squares” if the context is clear. So, for example, the numbers 1, 2, and 4 are squares mod 7, and the numbers 3, 5, and 6 are not.

If k is a square mod p we will call is a low square if 0 ≤ k < p/2 and a high square if p/2 < k < p.

Signatures

Now let p > 3 be a prime congruent to 3 mod 4. Add up all the low squares mod p and take the remainder mod p. Call this the signature of p. Here’s Python code to make this explicit.

from sympy import isprime, factorint, is_quad_residue

def signature(p):
    assert(p > 3)
    assert(isprime(p))
    assert(p % 4 == 3)
    s = 0
    for k in range(1, 1 + p//2):
        if is_quad_residue(k, p):
            s += k
    return s % p

Inverse signatures

Surprisingly, the signature of each p is unique. Given the signature of p, you can uniquely determine p, and in fact you can do so easily. I ran across this in a paper [1] that presented the results in the form of a parlor trick: have someone pick a prime p such that p = 3 mod 4 and ask them to compute its signature, the sum of the low squares mod p. Then you can quickly tell them what their choice of p was.

Given a signature s, the corresponding prime p is the largest prime factor of 16s + 1.

Not only that,

p = (16s + 1)/m

where m is the smallest of the numbers {3, 7, 11, 15} such that the fraction above is a prime number. In term of Python code, both the following functions should invert the signature of p.

def inverse_signature1(s):
    n = 16*s + 1
    return max(factorint(n).keys())

def inverse_signature2(s):
    n = 16*s + 1
    for m in [3, 7, 11, 15]:
        if n % m == 0 and isprime(n // m):
            return n // m

The following code demonstrates that this is the case for numbers less than 1,000.

for n in range(7, 1000, 4):
    if isprime(n):
        s = signature(n)
        assert(n == inverse_signature1(s))
        assert(n == inverse_signature2(s))

[1] David M. Bloom. A Quadratic Residues Parlor Trick. Mathematics Magazine, Vol. 71, No. 3 (Jun., 1998), pp. 201–203.

The post Sum of low squares first appeared on John D. Cook.

John D. Cook Modify

Rss preview of Blog of John D. Cook

More forensics posts

Solutions

SWI Prolog code

Lean 4 code

Beating brute force

Prefixes versus suffixes

HCPCS examples

Squares, high and low

Signatures

Inverse signatures

John D. CookModify

Rss preview of Blog of John D. Cook

More forensics posts

Solutions

SWI Prolog code

Lean 4 code

Beating brute force

Prefixes versus suffixes

HCPCS examples

Squares, high and low

Signatures

Inverse signatures

The author's social media

John D. Cook Modify