Problem 451

completed March 20, 2021

Comments

First day of spring! Happy mathing and happy sunshine!

I get frustrated with myself when I can’t figure out a problem. I honestly spent so many hours studying roots of unity modulo n and modular arithmetic in general. At least two days this past week after work I spent all evening (at least 5 hours each night) trying to figure it out. No luck. I kept getting closer and closer with each try.

Instead, I decided to turn to a different approach. Instead of throwing math at the problem, I just threw a bunch of CPUs at it. I figure the laptop I’m working on has 16 cores and Python – being just a single process – will only ever utilize one. Golang on the other hand has goroutines which by default will take advantage of all its available CPUs. This allowed me to do computations 16 times faster.

I created 16 workers (this number chosen dynamically based on available CPU) and then had each worker pull batches of numbers to process. Some batches take longer than others, so I kept the batch size relatively small which meant that all 16 CPU were being used the entire time.

My algorithm was pretty brute force. Given the equation \(b^2 \equiv 1 \pmod{n}\), for each possible value of \(b\) I went about calculating the available values of \(n\). This is possible because \(b^2 = mn + 1\) where \(m\) is an integer.

Reading through the forum afterward makes me even more frustrated with myself. I’m sure I read all about the theorems that others used to solve this problem. I guess I just wasn’t able to put it together. But, with time and practice, I’m sure I’ll be able to do it in the future.

Final Go code takes a couple hours to run.

EDIT: Looking back at the forums again, I see people referencing Problem 407 and even just using the code they wrote for that problem in this one. Well, I’ve gotta work on Problem 407 next then, I told myself, only to find out I’d already solved it… The first sentence in my write up from that problem was

I never really grokked this one, but I did end up getting the right answer.

Clearly. Just shows me that I should really be taking the time after a tough problem to go through the answers in the forum and really understand them. That knowledge will certainly pay off later. I’m gonna go back and do that now…

Python

Code

def solve(n):
    unity_roots = [0] * 3 + [1] * (n - 2)
    for b in range(n // 2 - 1, 2, -1):
        b2_1 = b**2 - 1
        for m in range((b - 2) // 2 + 1, 0, -1):
            num, mod = divmod(b2_1, m)
            if num > n:
                break
            if not mod:
                unity_roots[num] = num - b

    return sum(unity_roots)

if __name__ == '__main__':
    import sys
    n = eval(sys.argv[1])
    print(solve(n))

Tests

import pytest
from problem import solve

_test_solve = (
        (10**1, 12),
        (10**2, 2044),
        (10**3, 278340),
        (10**4, 32246295),
        (10**5, 3493155809),
        #(10**6, 367567257265),
)

@pytest.mark.parametrize('n,expect', _test_solve)
def test_solve(n, expect):
    assert expect == solve(n)

Go

Code

package main

import (
	"fmt"
	"log"
	"math"
	"os"
	"runtime"
	"strconv"
	"sync"
	"time"
)

func main() {
	n, err := strconv.Atoi(os.Args[1])
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("%#v", Problem(n))
}

type root struct {
	key   int
	value int
}

func divmod(n, d int) (q, r int) {
	q = n / d
	r = n % d
	return
}

type manager struct {
	sync.WaitGroup
	found       chan root
	done        chan struct{}
	ans         chan int
	batches     chan batch
	batchesDone chan struct{}
}

func newManager() *manager {
	return &manager{
		found:       make(chan root),
		done:        make(chan struct{}),
		ans:         make(chan int),
		batches:     make(chan batch),
		batchesDone: make(chan struct{}),
	}
}

func (m *manager) manage(n int) {
	go func() {
		unity := make([]int, n+1)
		for {
			select {
			case r := <-m.found:
				if r.value > unity[r.key] {
					unity[r.key] = r.value
				}
			case <-m.done:
				var total int
				for i, unit := range unity {
					if i < 3 {
						continue
					}
					if unit == 0 {
						total++
					} else {
						total += unit
					}
				}
				m.ans <- total
				return
			}
		}
	}()

	go func() {
		min := 3
		max := 3
		stop := n/2 - 1
		size := 1000
		var complete bool
		for {
			max = min + size
			if max > stop {
				max = stop
				complete = true
			}
			m.batches <- batch{min: min, max: max}
			min += size + 1
			if complete {
				close(m.batchesDone)
				return
			}
		}
	}()
}

func (m *manager) answer() int {
	m.Wait()
	close(m.done)
	return <-m.ans
}

type batch struct {
	min int
	max int
}

func (m *manager) process(n int, b batch) {
	for a := b.min; a <= b.max; a++ {
		b21 := int(math.Pow(float64(a), 2)) - 1
		for x := (a-2)/2 + 1; x > 0; x-- {
			num, mod := divmod(b21, x)
			if num > n {
				break
			}
			if mod == 0 {
				m.found <- root{key: num, value: num - a}
			}
		}
	}
}

func Problem(n int) int {
	m := newManager()
	go m.manage(n)

	cpus := runtime.NumCPU()
	for c := 0; c < cpus; c++ {
		time.Sleep(5 * time.Millisecond)
		m.Add(1)
		go func() {
			for {
				select {
				case batch := <-m.batches:
					m.process(n, batch)
				case <-m.batchesDone:
					m.Done()
					return
				}
			}
		}()
	}
	return m.answer()
}

Tests

package main

import (
	"fmt"
	"math"
	"testing"
)

func TestProblem(t *testing.T) {
	testcases := []struct {
		n      int
		expect int
	}{
		{n: int(math.Pow(10, 1)), expect: 12},
		{n: int(math.Pow(10, 2)), expect: 2044},
		{n: int(math.Pow(10, 3)), expect: 278340},
		{n: int(math.Pow(10, 4)), expect: 32246295},
		{n: int(math.Pow(10, 5)), expect: 3493155809},
		//{n: int(math.Pow(10, 6)), expect: 367567257265},
	}

	for _, test := range testcases {
		t.Run(fmt.Sprintf("n: %#v\n", test.n), func(t *testing.T) {
			if actual := Problem(test.n); actual != test.expect {
				t.Errorf("wrong answer for n=%d: actual=%d expect=%d",
					test.n, actual, test.expect)
			}
		})
	}
}