Partitions in Arrays with at least K-Size and M-Difference Rules (Array Partition)

Given an array A[] of N integers, find if it is possible to partition the array with following rules:

• Each element should belong to exactly one partition.
• Each partition should have at least K elements.
• Absolute difference between any pair of elements in the same partition should not exceed M.

Examples:

Input: N = 5, K = 2, M = 3, A[] = {8, 3, 9, 1, 2}
Output: YES
Explanation: We can partition the array into two partitions: {8, 9} and {3, 1, 2} such that all rules are satisfied.

Input: N = 6, K = 2, M = 3, A[] = {4, 3, 8, 1, 2}
Output: NO
Explanation: No possible partition exists that follows all the given conditions.

Approach: The problem can be solved using the following approach:

The idea is to use a Segment tree to efficiently track valid partition ranges. Sort the input array so that the difference between the adjacent elements are as small as possible. Now, while iterating through the sorted array, identify the left and right endpoints of the valid partitions such the absolute difference <= M and the number of elements is at least K. The right most index is calculated using the absolute difference limit M and the left most index is calculated using the partition limit K.

Steps to solve the problem:

• Declare two global arrays: seg[] array to store information about partitioning possibilities and marked[] array to store whether a node in the segment tree is marked or not.
• Maintain a function push() to propagate changes in the segment tree. If a node is marked (marked[v] is true), the changes are propagated to its children (seg[2 * v] and seg[2 * v + 1]), changing the marked flag of the current node as false and marked flags of its children as true.
• Maintain a function get_segment() to retrieve the segment value at a particular position in the segment tree.
• Maintain another function update_segment() to update segment values in a given range.
• Iterate over the sorted array A[]:
  • Determine the rightmost index (right) within the absolute difference limit (M) using upper_bound().
  • Calculate the left index (left) based on the current index (i) and partition size (K).
  • The variable left represents the left index of the range to be considered for updating the segment tree.

• If it’s the first element, update the segment tree for the given range (update_segment).
  • If the previous element can be part of a partition (checked using get_segment), update the segment tree for the given range.
• Check if the last element can be part of a partition using get_segment.
• Output “YES” if possible; otherwise, output “NO”.

Below is the implementation of the above approach:

C++

#include <algorithm>
#include <iostream>
#include <vector>

using namespace std;

// declaring segment tree and marked arrays
vector<int> seg(800010);
vector<bool> marked(800010);

// function to push changes in segment tree
void push(int v)
{
    // if node is marked, propagate the changes to its
    // children
    if (marked[v]) {
        seg[2 * v] |= seg[v];
        seg[2 * v + 1] |= seg[v];
        marked[v] = false;
        marked[2 * v] = marked[2 * v + 1] = true;
    }
}

// function to get segment value at a particular position
int get_segment(int v, int l, int r, int pos)
{
    // if leaf node, return segment value
    if (l == r) {
        return seg[v];
    }
    else {
        // push changes and recursively search for the value
        // in the left or right subtree
        push(v);
        int m = (l + r) / 2;
        if (pos <= m) {
            return get_segment(2 * v, l, m, pos);
        }
        else {
            return get_segment(2 * v + 1, m + 1, r, pos);
        }
    }
}

// function to update segment values in a given range
void update_segment(int v, int l, int r, int query_l,
                    int query_r)
{
    // if range matches exactly with the node, update the
    // segment value and mark it
    if (l == query_l && r == query_r) {
        seg[v] |= 1;
        marked[v] = true;
    }
    else {
        // push changes and recursively update the left or
        // right subtree
        push(v);
        int m = (l + r) / 2;
        if (query_r <= m) {
            update_segment(2 * v, l, m, query_l, query_r);
        }
        else if (query_l > m) {
            update_segment(2 * v + 1, m + 1, r, query_l,
                           query_r);
        }
        else {
            update_segment(2 * v, l, m, query_l, m);
            update_segment(2 * v + 1, m + 1, r, m + 1,
                           query_r);
        }
        // update the segment value based on the values in
        // the left and right subtree
        seg[v] = seg[2 * v] | seg[2 * v + 1];
    }
}

int main()
{
    int N = 5;
    int K = 2;
    int M = 3;
    vector<int> A = { 8, 3, 9, 1, 2 };

    // initializing segment tree and marked array
    for (int i = 0; i <= 4 * N; i++) {
        seg[i] = 0;
        marked[i] = false;
    }

    // sorting the array in ascending order
    sort(A.begin(), A.end());

    // iterating over the array
    for (int i = 0; i < N; i++) {
        int right
            = int(upper_bound(A.begin(), A.end(), A[i] + M)
                  - A.begin());
        right--;
        int left = K + i - 1;

        // if first element, update the segment tree for the
        // given left and right range
        if (i == 0) {
            if (left <= right) {
                update_segment(1, 1, N, left + 1,
                               right + 1);
            }
        }
        // if the previous element can be partitioned,
        // update the segment tree for the given range
        else if (get_segment(1, 1, N, i)) {
            if (left <= right) {
                update_segment(1, 1, N, left + 1,
                               right + 1);
            }
        }
    }
    // check if the last element can be partitioned
    if (get_segment(1, 1, N, N)) {
        cout << "YES" << endl;
    }
    else {
        cout << "NO" << endl;
    }

    return 0;
}

Java

/*code by flutterfly */
import java.util.Arrays;

public class Main {

    // Declaring segment tree and marked arrays
    static int[] seg = new int[800010];
    static boolean[] marked = new boolean[800010];

    // Function to push changes in the segment tree
    static void push(int v) {
        // If the node is marked, propagate the changes to its children
        if (marked[v]) {
            seg[2 * v] |= seg[v];
            seg[2 * v + 1] |= seg[v];
            marked[v] = false;
            marked[2 * v] = marked[2 * v + 1] = true;
        }
    }

    // Function to get segment value at a particular position
    static int getSegment(int v, int l, int r, int pos) {
        // If leaf node, return segment value
        if (l == r) {
            return seg[v];
        } else {
            // Push changes and recursively search for the value
            // in the left or right subtree
            push(v);
            int m = (l + r) / 2;
            if (pos <= m) {
                return getSegment(2 * v, l, m, pos);
            } else {
                return getSegment(2 * v + 1, m + 1, r, pos);
            }
        }
    }

    // Function to update segment values in a given range
    static void updateSegment(int v, int l, int r, int query_l, int query_r) {
        // If range matches exactly with the node, update the
        // segment value and mark it
        if (l == query_l && r == query_r) {
            seg[v] |= 1;
            marked[v] = true;
        } else {
            // Push changes and recursively update the left or
            // right subtree
            push(v);
            int m = (l + r) / 2;
            if (query_r <= m) {
                updateSegment(2 * v, l, m, query_l, query_r);
            } else if (query_l > m) {
                updateSegment(2 * v + 1, m + 1, r, query_l, query_r);
            } else {
                updateSegment(2 * v, l, m, query_l, m);
                updateSegment(2 * v + 1, m + 1, r, m + 1, query_r);
            }
            // Update the segment value based on the values in
            // the left and right subtree
            seg[v] = seg[2 * v] | seg[2 * v + 1];
        }
    }

    public static void main(String[] args) {
        int N = 5;
        int K = 2;
        int M = 3;
        int[] A = {8, 3, 9, 1, 2};

        // Initializing segment tree and marked array
        Arrays.fill(seg, 0);
        Arrays.fill(marked, false);

        // Sorting the array in ascending order
        Arrays.sort(A);

        // Iterating over the array
        for (int i = 0; i < N; i++) {
            int right = Arrays.binarySearch(A, A[i] + M);
            right = right < 0 ? -right - 2 : right;
            int left = K + i - 1;

            // If the first element, update the segment tree for the
            // given left and right range
            if (i == 0) {
                if (left <= right) {
                    updateSegment(1, 1, N, left + 1, right + 1);
                }
            }
            // If the previous element can be partitioned,
            // update the segment tree for the given range
            else if (getSegment(1, 1, N, i) != 0) {
                if (left <= right) {
                    updateSegment(1, 1, N, left + 1, right + 1);
                }
            }
        }

        // Check if the last element can be partitioned
        if (getSegment(1, 1, N, N) != 0) {
            System.out.println("YES");
        } else {
            System.out.println("NO");
        }
    }
}

Python3

# code by flutterfly
from bisect import bisect_right

# Declaring segment tree and marked arrays
seg = [0] * 800010
marked = [False] * 800010

# Function to push changes in the segment tree
def push(v):
    # If node is marked, propagate the changes to its children
    if marked[v]:
        seg[2 * v] |= seg[v]
        seg[2 * v + 1] |= seg[v]
        marked[v] = False
        marked[2 * v] = marked[2 * v + 1] = True

# Function to get segment value at a particular position
def get_segment(v, l, r, pos):
    # If leaf node, return segment value
    if l == r:
        return seg[v]
    else:
        # Push changes and recursively search for the value
        # in the left or right subtree
        push(v)
        m = (l + r) // 2
        if pos <= m:
            return get_segment(2 * v, l, m, pos)
        else:
            return get_segment(2 * v + 1, m + 1, r, pos)

# Function to update segment values in a given range
def update_segment(v, l, r, query_l, query_r):
    # If range matches exactly with the node, update the
    # segment value and mark it
    if l == query_l and r == query_r:
        seg[v] |= 1
        marked[v] = True
    else:
        # Push changes and recursively update the left or
        # right subtree
        push(v)
        m = (l + r) // 2
        if query_r <= m:
            update_segment(2 * v, l, m, query_l, query_r)
        elif query_l > m:
            update_segment(2 * v + 1, m + 1, r, query_l, query_r)
        else:
            update_segment(2 * v, l, m, query_l, m)
            update_segment(2 * v + 1, m + 1, r, m + 1, query_r)
        # Update the segment value based on the values in
        # the left and right subtree
        seg[v] = seg[2 * v] | seg[2 * v + 1]

def main():
    N = 5
    K = 2
    M = 3
    A = [8, 3, 9, 1, 2]

    # Initializing segment tree and marked array
    for i in range(4 * N + 1):
        seg[i] = 0
        marked[i] = False

    # Sorting the array in ascending order
    A.sort()

    # Iterating over the array
    for i in range(N):
        right = bisect_right(A, A[i] + M) - 1
        left = K + i - 1

        # If the first element, update the segment tree for the
        # given left and right range
        if i == 0:
            if left <= right:
                update_segment(1, 1, N, left + 1, right + 1)
        # If the previous element can be partitioned,
        # update the segment tree for the given range
        elif get_segment(1, 1, N, i) != 0:
            if left <= right:
                update_segment(1, 1, N, left + 1, right + 1)

    # Check if the last element can be partitioned
    if get_segment(1, 1, N, N) != 0:
        print("YES")
    else:
        print("NO")

if __name__ == "__main__":
    main()

C#

//code by flutterfly
using System;
using System.Collections.Generic;

class Program
{
    static List<int> seg = new List<int>();
    static List<bool> marked = new List<bool>();

    static void Push(int v)
    {
        // If the node is marked, propagate the changes to its children
        if (marked[v])
        {
            seg[2 * v] |= seg[v];
            seg[2 * v + 1] |= seg[v];
            marked[v] = false;
            marked[2 * v] = marked[2 * v + 1] = true;
        }
    }

    static int GetSegment(int v, int l, int r, int pos)
    {
        // If leaf node, return segment value
        if (l == r)
        {
            return seg[v];
        }
        else
        {
            // Push changes and recursively search for the value
            // in the left or right subtree
            Push(v);
            int m = (l + r) / 2;
            if (pos <= m)
            {
                return GetSegment(2 * v, l, m, pos);
            }
            else
            {
                return GetSegment(2 * v + 1, m + 1, r, pos);
            }
        }
    }

    static void UpdateSegment(int v, int l, int r, int query_l, int query_r)
    {
        // If range matches exactly with the node, update the
        // segment value and mark it
        if (l == query_l && r == query_r)
        {
            seg[v] |= 1;
            marked[v] = true;
        }
        else
        {
            // Push changes and recursively update the left or
            // right subtree
            Push(v);
            int m = (l + r) / 2;
            if (query_r <= m)
            {
                UpdateSegment(2 * v, l, m, query_l, query_r);
            }
            else if (query_l > m)
            {
                UpdateSegment(2 * v + 1, m + 1, r, query_l, query_r);
            }
            else
            {
                UpdateSegment(2 * v, l, m, query_l, m);
                UpdateSegment(2 * v + 1, m + 1, r, m + 1, query_r);
            }
            // Update the segment value based on the values in
            // the left and right subtree
            seg[v] = seg[2 * v] | seg[2 * v + 1];
        }
    }

    static void Main()
    {
        int N = 5;
        int K = 2;
        int M = 3;
        List<int> A = new List<int> { 8, 3, 9, 1, 2 };

        // Initializing segment tree and marked array
        seg = new List<int>(new int[4 * N + 1]);
        marked = new List<bool>(new bool[4 * N + 1]);

        // Sorting the array in ascending order
        A.Sort();

        // Iterating over the array
        for (int i = 0; i < N; i++)
        {
            int right = A.BinarySearch(A[i] + M);
            right = (right < 0) ? ~right - 1 : right;
            int left = K + i - 1;

            // If the first element, update the segment tree for the
            // given left and right range
            if (i == 0)
            {
                if (left <= right)
                {
                    UpdateSegment(1, 1, N, left + 1, right + 1);
                }
            }
            // If the previous element can be partitioned,
            // update the segment tree for the given range
            else if (GetSegment(1, 1, N, i) != 0)
            {
                if (left <= right)
                {
                    UpdateSegment(1, 1, N, left + 1, right + 1);
                }
            }
        }

        // Check if the last element can be partitioned
        if (GetSegment(1, 1, N, N) != 0)
        {
            Console.WriteLine("YES");
        }
        else
        {
            Console.WriteLine("NO");
        }
    }
}

Javascript

// Javascript program for the above approach
const upper_bound = (arr, target) => {
  let left = 0;
  let right = arr.length;

  while (left < right) {
    const mid = Math.floor((left + right) / 2);
    if (arr[mid] <= target) {
      left = mid + 1;
    } else {
      right = mid;
    }
  }

  return left;
};

// declaring segment tree and marked arrays
const seg = new Array(800010);
const marked = new Array(800010);

// function to push changes in segment tree
function push(v) {
  // if node is marked, propagate the changes to its
  // children
  if (marked[v]) {
    seg[2 * v] |= seg[v];
    seg[2 * v + 1] |= seg[v];
    marked[v] = false;
    marked[2 * v] = marked[2 * v + 1] = true;
  }
}

// function to get segment value at a particular position
function get_segment(v, l, r, pos) {
  // if leaf node, return segment value
  if (l === r) {
    return seg[v];
  } else {
    // push changes and recursively search for the value
    // in the left or right subtree
    push(v);
    const m = Math.floor((l + r) / 2);
    if (pos <= m) {
      return get_segment(2 * v, l, m, pos);
    } else {
      return get_segment(2 * v + 1, m + 1, r, pos);
    }
  }
}

// function to update segment values in a given range
function update_segment(v, l, r, query_l, query_r) {
  // if range matches exactly with the node, update the
  // segment value and mark it
  if (l === query_l && r === query_r) {
    seg[v] |= 1;
    marked[v] = true;
  } else {
    // push changes and recursively update the left or
    // right subtree
    push(v);
    const m = Math.floor((l + r) / 2);
    if (query_r <= m) {
      update_segment(2 * v, l, m, query_l, query_r);
    } else if (query_l > m) {
      update_segment(2 * v + 1, m + 1, r, query_l, query_r);
    } else {
      update_segment(2 * v, l, m, query_l, m);
      update_segment(2 * v + 1, m + 1, r, m + 1, query_r);
    }
    // update the segment value based on the values in
    // the left and right subtree
    seg[v] = seg[2 * v] | seg[2 * v + 1];
  }
}

const N = 5;
const K = 2;
const M = 3;
const A = [8, 3, 9, 1, 2];

// initializing segment tree and marked array
for (let i = 0; i <= 4 * N; i++) {
  seg[i] = 0;
  marked[i] = false;
}

// sorting the array in ascending order
A.sort((a, b) => a - b);

// iterating over the array
for (let i = 0; i < N; i++) {
  const right = upper_bound(A, A[i] + M) - 1;
  const left = K + i - 1;

  // if first element, update the segment tree for the
  // given left and right range
  if (i === 0) {
    if (left <= right) {
      update_segment(1, 1, N, left + 1, right + 1);
    }
  }
  // if the previous element can be partitioned,
  // update the segment tree for the given range
  else if (get_segment(1, 1, N, i)) {
    if (left <= right) {
      update_segment(1, 1, N, left + 1, right + 1);
    }
  }
}
// check if the last element can be partitioned
if (get_segment(1, 1, N, N)) {
  console.log("YES");
}
else {
  console.log("NO");
}

// This code is contributed by Susobhan Akhuli

Output

YES

Time Complexity: O(N log(N)), where N is the number of elements in the input array A[].
Auxiliary Space: O(N)