# Count switches in a smallest square root sequence mod $2^n$

### Definition

For any $$\a\equiv1\ (\text{mod }8)\$$ and $$\n\ge3\$$, there are exactly 4 roots to the equation $$\x^2\equiv a\ (\text{mod }2^n)\$$. Now, let $$\x_k(a)\$$ be the smallest root to the equation $$\x^2\equiv a\ (\text{mod }2^k)\$$, then $$\{x_3(a),x_4(a),x_5(a),x_6(a),\cdots\}$$ is a smallest square root sequence (SSRS) of $$\a\$$ mod $$\2^n\$$.

John D. Cook published a quick algorithm that calculates such roots in $$\O(n)\$$ time. Assume $$\x_k\$$ is a root to the equation $$\x^2\equiv a\ (\text{mod }2^k)\$$. Then, $$x_{k+1}=\begin{cases}x_k&\text{if }\frac{x_k^2-a}{2^k}\text{ is even}\\x_k+2^{k-1}&\text{otherwise}\end{cases}$$ is a root to the equation $$\x^2\equiv a\ (\text{mod }2^{k+1})\$$.

Now we define two lists A and B. $$\A=\{A_k|k\ge3\}\$$ is the list of values generated by the algorithm above with initial values $$\A_3=1\$$ and $$\B=\{B_k|k\ge3\}\$$ is the list of values generated with initial values $$\B_3=3\$$. Each entry in the SSRS $$\x_k(a)\$$ takes the smallest value among $$\A_k\$$ and $$\B_k\$$. We say a switch in SSRS occurs whenever the choice changes from A to B or from B to A.

To illustrate the definition, take $$\a=17\$$: The smallest numbers are highlighted. From the picture there are 13 switches up to mod $$\2^{24}\$$.

### Challenge

Write a function or program, that receives 2 integers $$\a,\ k\$$ as input (where $$\a\equiv1\ (\text{mod }8)\$$ and $$\k\ge3\$$) and output how many switches occur in the SSRS of $$\a\$$ mod $$\2^n\$$ up to $$\n=k\$$.

### Sample I/O

1, 3 -> 0
9, 4 -> 1
1, 8 -> 0
9, 16 -> 1
17, 24 -> 13
25, 32 -> 2
33, 40 -> 18
41, 48 -> 17
49, 56 -> 1
1048577, 2048 -> 959
1048585, 2048 -> 970


### Winning Condition

This is a code-golf challenge, so shortest valid submission of each language wins. Standard loopholes are forbidden by default.

• Getting 1048577, 2048 --> 959 with all the other test cases correct. Am I missing something? – Noodle9 Feb 10 at 13:13
• Was this intended to be restricted-complexity? Otherwise we can just ignore Cook's formula and brute-force the square roots. – Grimmy Feb 10 at 13:13
• @Noodle9 It is 959 indeed! I have amended the test cases. – Shieru Asakoto Feb 10 at 13:48
• @Grimmy If bruteforcing gives the correct results I will accept it. – Shieru Asakoto Feb 10 at 13:50
• Oh I finally know why I put 1018 in the (1048577, 2048) case. It's actually the result of (17, 2048). – Shieru Asakoto Feb 11 at 3:52

# 05AB1E, 15 14 bytes

LoεLnαy%0k}üÊO


Try it online!

For each n, this finds the smallest square root of a by brute-force.

## 05AB1E, 26 23 bytes (fast)

13v¹yλ£DnIαN>o%2÷+]@γg<


Try it online!

This one properly uses Cook's formula.

13                    # literal 13
v               ]   # for each digit:
¹                  #  push the input k
y             ]   #  push the digit
λ£               #  recurse k times with base case y:
Dn             #   square of the current value
Iα           #   absolute difference with input a
N>o%       #   modulo 2**(N+1)
2÷     #   integer divide by 2
+    #   add to the current value
@                     # compare the two lists element-wise
γ                    # group consecutive equal elements
g                   # length
<                  # -1


# JavaScript (ES6),  88  87 bytes

Takes input as (a)(k).

Because of the bitwise operations, this is only guaranteed to work for $$\k\le32\$$.

a=>K=>(x=[p=1,3],F=k=>k<K&&(p^(x=x.map(v=>q=v+=(v*v-a>>k&1)<<k-1),p=x<q))+F(k+1))(2)


Try it online!

# Charcoal, 50 bytes

≔⁰ζ⊞υ¹⊞υ³ＦＸ²…³Ｎ«ＵＭυ⁺κ∧﹪÷⁻×κκηι²⊘ιＦ⌕υ⌊υ«≦⊕ζ≔⮌υυ»»Ｉζ


Try it online! Link is to verbose version of code. Takes input in the order $$\ \small k, a \$$. Explanation:

≔⁰ζ⊞υ¹⊞υ³


Start off with $$\ \small 0 \$$ swaps and initial values $$\ \small A_3 = 1 \$$ and $$\ \small B_3 = 3 \$$ in a list.

ＦＸ²…³Ｎ«


Loop over the powers of $$\ \small 2^n \$$ from $$\ \small 2^3 \$$ to $$\ \small 2^{k-1} \$$.

ＵＭυ⁺κ∧﹪÷⁻×κκηι²⊘ι


Map over the list calculating the next values $$\ \small A_{n+1} \$$ and $$\ \small B_{n+1} \$$.

Ｆ⌕υ⌊υ«


Test whether the smallest value is at the start of the list.

≦⊕ζ≔⮌υυ


If it isn't then increment the number of swaps and reverse the list so that the smallest value is at the start again.

»»Ｉζ


Print the number of swaps.

I have a 48 47 byte version that seems to work on the test cases but I don't know why.

≔⁰ζ≔¹δ≔⁰εＦＸ²…³Ｎ«≔﹪÷⁻Ｘδ²ηι²θ≧⁺¬⁼θεζ≔θε≧⁺∧ε⊘ιδ»Ｉζ


Try it online! Link is to verbose version of code. Explanation:

≔⁰ζ≔¹δ≔⁰ε


Start off with $$\ \small 0 \$$ swaps, $$\ \small A_3 = 1 \$$ and $$\ \small A_3 < B_3 \$$.

ＦＸ²…³Ｎ«


Loop over the powers of $$\ \small 2^n \$$ from $$\ \small 2^3 \$$ to $$\ \small 2^{k-1} \$$.

≔﹪÷⁻Ｘδ²ηι²θ


Work out whether $$\ \small A_{n+1} > B_{n+1} \$$.

≧⁺¬⁼θεζ


If this is a change from $$\ \small A_n > B_n \$$ then increment the number of swaps.

≔θε


Save this for the next loop.

≧⁺∧ε⊘ιδ


Add $$\ \small 2^{n-1} \$$ if $$\ \small A_{n+1} > B_{n+1} \$$.

»Ｉζ


Print the number of swaps.

# Jelly, 31 29 bytes

²_⁴ọ2>
R1,3;€¹2*H+ʋ@ç?\€>/IAS


Try it online!

A full program taking $$\n\$$ as its left argument and $$\k\$$ as its right argument. Prints the number of switches.

Will post explanation once have more time to further optimise.

# Python 3, 230 $$\\cdots\$$ 180 175 bytes

def f(x,a,k,p=4):
l=[x]
for n in range(3,k):x+=p*(((x*x-a)>>n)&1);p*=2;l+=[x]
return l
def g(a,n):
i=c=0
for q in zip(f(1,a,n),f(3,a,n)):j=q[1-i]<q[i];c+=j;i^=j
return c


Try it online!

A function g that uses Cook's formula (function f).