Calculate the integer square root of a matrix

Question

Let \$A\$ be a square matrix that is at least \$2 \times 2\$ where each element is an integer. \$A^2 = A \times A\$ will then have the same dimensions as \$A\$, and will have integer elements. For example, let

$$A = \left[ \begin{matrix} -3 & 2 \\ 0 & -1 \end{matrix} \right]$$

Therefore,

$$\begin{align} A^2 & = \left[ \begin{matrix} -3 & 2 \\ 0 & -1 \end{matrix} \right]^2 \\ & = \left[ \begin{matrix} -3 & 2 \\ 0 & -1 \end{matrix} \right] \times \left[ \begin{matrix} -3 & 2 \\ 0 & -1 \end{matrix} \right] \\ & = \left[ \begin{matrix} -3 \times -3 + 2 \times 0 & -3 \times 2 + 2 \times -1 \\ 0 \times -3 + -1 \times 0 & 0 \times 2 + -1 \times -1 \\ \end{matrix} \right] \\ & = \left[ \begin{matrix} 9 & -8 \\ 0 & 1 \end{matrix} \right] \end{align}$$

---

Given an \$n \times n\$ matrix \$B\$ consisting only of integer elements, output any \$n \times n\$ integer matrix \$A\$ such that \$A^2 = B\$. You may optionally take \$n\$ as an input, and you may assume that you'll only have to handle matrices for which there exists such an \$A\$.

You may take input and output in any convenient format and method. This is code-golf so the shortest code in bytes wins.

Test cases

_{The outputs given below are not necessarily the only valid outputs}

[[8, 96], [-84, -40]] -> [[-8, -8], [7, -4]]
[[18, 63], [14, 67]] -> [[0, 9], [2, 7]]
[[9, 0], [-10, 25]] -> [[3, 0], [5, -5]]
[[25, -58, 57], [0, 7, -4], [0, -24, 31]] -> [[5, -7, 5], [0, -1, -1], [0, -6, 5]]
[[12, -33, 42], [16, 19, 20], [-90, 18, 63]] -> [[6, -3, 3], [2, 5, 1], [-6, 0, 9]]
[[33, -80, 86], [72, 57, -13], [-88, 40, 44]] -> [[-8, 7, -3], [-1, 0, 8], [8, 8, -2]]
[[88, 8, -72, 65], [-12, 45, 17, 33], [-27, 21, 36, 31], [5, 40, -53, 119]] -> [[-5, -4, 3, -9], [-6, 0, 0, 7], [-5, 5, -4, 7], [-6, 3, 5, -3]]
[[45, 39, -25, -17, 61], [29, 69, -15, 2, 45], [42, 51, 7, -28, 67], [27, 65, -25, 7, 25], [-66, -61, 9, 63, 1]] -> [[9, 7, -3, -9, 3], [3, -3, 5, -8, -2], [5, 3, 5, -9, 2], [3, -2, 2, -8, -4], [-5, -4, 0, 7, 6]]
[[150, -73, -37, -40, -43, 119], [-62, 191, 95, -87, -10, -88], [-31, -32, -64, 137, 82, -54], [22, -81, 32, 24, 46, -149], [-5, -16, -6, 42, 63, -23], [106, -160, -115, 25, 20, -5]] -> [[8, -8, -3, 5, 6, 3], [-7, 6, 8, 5, 6, -4], [5, 1, -1, -7, 1, -9], [0, 9, 6, -8, -2, 9], [4, 3, 2, -1, -3, -2], [7, -7, 2, -5, -4, 1]]
[[146, -37, -60, -38, 30, -8], [71, -42, -33, -26, -15, -55], [-39, -74, 2, -46, -19, 31], [42, -136, -46, 19, -41, -3], [61, 23, 12, -8, 37, -59], [-73, -13, 88, 90, -71, -11]] -> [[4, 7, -2, -4, 4, -1], [8, -1, -3, -6, 2, 4], [-9, 5, -2, 6, 1, 3], [-7, 6, 1, -5, 4, 3], [5, -8, -8, 2, 3, -4], [-8, -8, -1, -4, -6, -2]]

Answer 1

Haskell, ⁹⁸ 96 bytes

• -2 bytes thanks to Bubbler.

s a=[x|m<-[1..],x<-q(q[-m..m]a)a,a==[foldr1(z(+))$z(map.(*))r x|r<-x]]!!0
q=mapM.const
z=zipWith

Try it online!

The relevant function is s, which takes as input a matrix a as a list of lists of integers and returns a square root of a in the same form.

The matrix multiplication part is taken almost verbatim from xnor's answer to this challenge.

Answer 2

WolframLanguage (Mathematica), 54 52 45 bytes

Last@Solve[(q=s~Array~{#,#}).q==#2,Integers]&

–2 bytes from ovs.

–7 bytes due to reading the directions and realizing I can take n as an input.

Try it online! [The request exceeded the 60 sec time limit, so it could only solve the first eight of the ten test cases.]

How it works:

q=s~Array~{#,#} //creates a dummy-indexed matrix with dimensions equal to *n* , calls it "q"
Solve[(q=s~Array~{#,#}).q==#2,Integers] //Solves q.q=#2 for q, where #2 is the input matrix, restricting the solutions to matrices with integer elements
Last@... //since there can be multiple solutions, this selects  the last one (the first ones contain conditionals)

Here I leaned heavily on the OP's allowance that output could be in "any convenient format", since the above code specifies the solution matrix in terms of its indexed elements, i.e., s[row, column]. For example:

For 3 additional bytes (48 total), a nicer output can be obtained with a small modification of the code offered by att in the comments:

(q=s~Array~{#,#})/.Last@Solve[q.q==#2,Integers]&

Answer 3

Julia 1.0, 109 100 bytes

I=Iterators
m%n=for i=I.countfrom(),j=I.product(fill(-i:i,n^2)...)
r=+m;r[:].=j;r^2==m&&return r end

Try it online!

9 bytes golfed by MarcMush.

Boring bruteforce solution. Takes input as matrix m and its size n. But here's also a bonus:

Julia 0.7, 214 194 bytes

~=round
m%n=.~filter(j->j≈.~j,1-2digits.(0:2^n-1,2,n).|>h->((t,z)=schur(m+0im);u=√t.*h;for i=n:-1:1,j=i+1:n a&k=a-u[i,k]*u[k,j];u[i,j]=reduce(&,t[i,j],i+1:j-1)/(u[i,i]+u[j,j])end;z)*u*z')[1]

Try it online!

This is an attempt at an algebraic solution making use of Julia's rich linear algebra library. In fact, Julia's sqrt or √ already works on matrices, but obviously only returns the principal root which may not be the one we need. Therefore, we resort to Schur decomposition, and our task boils down to finding the square roots of the upper triangular matrix \$T\$ (notation is different from Wiki, but corresponds to Julia's):

\$M^\frac{1}{2}=ZT^\frac{1}{2}Z^{-1}\$

The square roots \$U\$ of matrix \$T\$ can be found by Björck-Hammarling method, which is what Julia itself uses in the implementation of sqrt:

\$U_{ii}=T_{ii}^\frac{1}{2}\$

\$U_{ij}=\frac{T_{ij} - \sum_{k=i+1}^{j-1} U_{ik}U_{kj}}{U_{ii}+U_{jj}}, j>i\$

In order to find all solutions, we also need to apply all possible permutations of the signs of the roots in the first equation. The upper triangle can then be fixed according to the second equation.

All these computations introduce some roundoff errors, but when checking for integers, we can resolve this by using inexact comparison operator ≈. There is a problem with zero, for which it behaves the same as exact equality and comparisons fail. (It works just fine when applied to entire matrix that has both zeros and non-zeros).

The results are composed of complex numbers with zero imaginary part, so in TIO footer, we extract the real part for nicer display. It would be -2 bytes if it's acceptable to output the values without cleaning the roundoff errors, and further -3, if we can actually output all found integer solutions instead of just one.

Answer 4

Jelly, 17 bytes

FAṀŒRṗLƲṁ€ðæ*2iị⁸

Try it online!

It's been over a week, so I figured I'd reveal my brute-force Jelly approach. This makes the following assumption, which I haven't managed to prove, but can't find a counter example:

Let \$B\$ be the input \$n \times n\$ matrix, and \$A\$ be an \$n \times n\$ matrix such that \$A^2 = B\$

Let \$b = \max\{|B_{ij}|, 1 \le i,j \le n\}\$ i.e. the largest absolute element of \$B\$. This answer assumes that, assuming \$A\$ exists, there is at least one matrix \$A\$ such that all its elements are in the inclusive range \$[-b, b]\$

For example, for $$B = \left[ \begin{matrix} 9 & -8 \\ 0 & 1 \end{matrix} \right],$$ \$b = 9\$ and there exists at least one \$A\$ such that all elements are between \$-9\$ and \$9\$ (the example in the question, for example)

How it works

FAṀŒRṗLƲṁ€ðæ*2iị⁸ - Main link. Takes B on the left
F                 - Flatten B
 A                - Absolute values of each.
       Ʋ          - Last 4 links as a monad f(abs(flat(B)):
  Ṁ               -   Maximum
   ŒR             -   Bounced range; [-max, +max]
      L           -   Length i.e. number of elements of B
     ṗ            -   Powerset
         €        - Over each list:
        ṁ         -   Mold it like B
                    Call this list of matrices M
          ð       - Begin a new dyadic chain with M on the left and B on the right
           æ*2    - Matrix square of each matrix in M
              i   - Index of B in M
                ⁸ - Yield M
               ị  - Index back into M

One byte longer is a slightly more conventional way of extracting the root:

FAṀŒRṗLƲṁ€æ*2⁼ɗƇ⁸Ḣ

Try it online!

Here, the dyadic filter æ*2⁼ɗƇ means that it would save a byte over the dyadic chaining with the indexing. However, as we should only output one solution, and the Ḣ would chain to the filter without either ⁸ or ¹, the indexing here actually saves a byte.