Partitioning unit square with equal frequency rectangles

Question

If I had to partition the unit square $[0,1]\times[0,1]$ into $k^2$ rectangles such that the sum of their diagonals is minimum possible, I would simply choose the $k \times k$ grid of squares. Now suppose we also have a collection of $nk^2$ points in general position inside the unit square and impose the additional requirement that each rectangle in the partition should contain exactly $n$ points.

1. Is there an efficient algorithm to determine an optimal/near-optimal partition?

2. What if we slightly relax the assumption that each rectangle should contain an equal number of points?(as is necessary when the total number of points is not divisible by $k^2$)

This question is related to a procedure called data discretization in the field of data mining and statistics. See below.

[1] James Dougherty, Ron Kohavi, Mehran Sahami, Supervised and Unsupervised Discretization of Continuous Features, Machine Learning Proceedings 1995

Answer 1

Here is an answer inspired by redistricting and the shortest split line algorithm.

For any rectangle with $mn$ points, consider the $2m-2$ ways of dividing it horizontally or vertically into two rectangles with an integral multiple of $n$ points in each. Among these possible divisions of the rectangle into $R$ and $R’$, we can choose the one for which $E[R]+E[R’]$ is minimal, where $E$ is an appropriate scoring function.

Now apply this technique to divide the unit square into two rectangles, and apply it recursively to each of the rectangles that result.

For the scoring function of a rectangle $I\times J$ with $mn$ points, we can take $E[I\times J]=\sqrt{A^2 m^2/i^2+B^2 i^2}$ where $A=|I|$, $B=|J|$, $i=\max(1,\min(m,\sqrt{mA/B}))$. This gives an optimal sum of diagonal lengths if the distribution of points is uniform and $i$ is integral. Then the scoring is quick, and the total algorithm requires a number of steps which is $O(k^4)$.

Answer 2

Matt's algorithm works pretty well, while the modification I had suggested has slightly worse performance. Sharing some plots here (all using the scoring function Matt had suggested).

import numpy as np
import matplotlib.pyplot as plt

def score_rectange(As,Bs,m):
    if MODE == 1:
        A = np.sum(As)
        B = np.sum(Bs)
        i = np.maximum(1,np.minimum(m,(m*A/B)**0.5))
        score = ((A*m/i)**2+(B*i)**2)**0.5
        
    elif MODE == 2:
        A = (np.sum(As**2))**(0.5)
        B = (np.sum(Bs**2))**(0.5)
        
        i = np.maximum(1,np.minimum(m,(m*A/B)**0.5))                          
        score = ((A*m/i)**2+(B*i)**2)

    return(score)

def get_splits(z,m,boundz):
    n = int(len(z)/m)
    beg_idx = np.arange(0,np.round(m-1)*n+1,n).astype(int)
    end_idx = beg_idx+n-1
    beg_val = z[beg_idx]
    end_val = z[end_idx]
    
    #splits are defined to be exaclty half-way in between : might not be optimal
    split_cuts  = (end_val[:-1]+beg_val[1:])/2
    split_cuts  = np.concatenate(([boundz[0]],split_cuts,[boundz[1]]))
    split_lens  = np.diff(split_cuts)
    
    return split_cuts,split_lens

def find_best_split(lens_c,lens_o,m):    
    scores = np.array([score_rectange(lens_c[:i],lens_o,i)+score_rectange(lens_c[i:],lens_o,m-i) for i in range(1,len(lens_c))])
    best_split = np.argmin(scores)
    return(best_split+1,scores[best_split])

def split_rectangle(points,bounds,m):
    if m==1:
        return [bounds]
        
    x_s,y_s = list(zip(*points))
    x_s = np.array(x_s)
    y_s = np.array(y_s)
    
    ix_s = np.argsort(x_s)
    iy_s = np.argsort(y_s)
    
    cuts_x,lens_x = get_splits(x_s[ix_s],m,bounds[0])
    cuts_y,lens_y = get_splits(y_s[iy_s],m,bounds[1])
    
    x_bsplit_i,x_bsplit_s = find_best_split(lens_x,lens_y,m)
    y_bsplit_i,y_bsplit_s = find_best_split(lens_y,lens_x,m)
    
    if x_bsplit_s<=y_bsplit_s:
        best_cut = cuts_x[x_bsplit_i]
        mask = x_s <= best_cut
        
        newm = (x_bsplit_i,m-x_bsplit_i)
        
        newbounds = (
            ((bounds[0][0],best_cut),bounds[1]),
            ((best_cut,bounds[0][1]),bounds[1])
        )
    else:
        best_cut = cuts_y[y_bsplit_i]
        mask = y_s <= best_cut
        
        newm = (y_bsplit_i,m-y_bsplit_i)
        
        newbounds = (
            (bounds[0],(bounds[1][0],best_cut)),
            (bounds[0],(best_cut,bounds[1][1]))
        )
    
    return(
        split_rectangle(list(zip(x_s[mask],y_s[mask])),newbounds[0],newm[0])+
        split_rectangle(list(zip(x_s[~mask],y_s[~mask])),newbounds[1],newm[1])
    )
    
    
x = np.random.beta(1,1,size=50000)
y = np.random.beta(1,1,size=50000)
y = x*0.7+y*0.3

points = list(zip(x,y))

MODE = 1
splits = split_rectangle(points,((0,1),(0,1)),25)
fig,ax = plt.subplots(1,1,figsize=(7,7))

sum_diag = 0
sum_diag2 = 0
for p in splits:
    x0 = p[0][0]
    x1 = p[0][1]
    y0 = p[1][0]
    y1 = p[1][1]
    
    diag2 = (x1-x0)**2+(y1-y0)**2
    diag = diag2**0.5
    sum_diag += diag
    sum_diag2 += diag2
    
    ax.vlines(x0,y0,y1)
    ax.vlines(x1,y0,y1)
    ax.hlines(y0,x0,x1)
    ax.hlines(y1,x0,x1)
    ax.scatter(x,y,color='grey',alpha=0.1)

ax.set_title("Sum of Diagonals = {:.2f} \n Sum of Squares of Diagonals = {:.2f}".format(sum_diag,sum_diag2));