CSES Solutions – Two Sets

Given N numbers 1,2,3,… N. Your task is to divide the numbers into two sets of equal sum. Print “YES“, if the division is possible, and “NO” otherwise. If the division is possible, print how to create the sets. First, print the number of elements in the first set followed by the elements themselves in a separate line, and then, similarly print the second set.

Examples:

Input: 7
Output:
YES
4
1 2 4 7
3
3 5 6
Explanation: The first set contains 4 elements which are 1,2,4 and 7 with sum equal to 14 and the second set contains 3 elements 3,5 and 6 with sum equal to 14.

Input: 6
Output: NO
Explanation: There is no possible way to divide numbers 1, 2, 3, 4, 5, 6 into 2 sets with equal sums.

Approach: To solve the problem, follow the below idea:

The idea is to find the total sum of n numbers. If total sum is odd, then no answer is possible else the answer always exists. Now we have to make the sum of one of the sets half of the total sum, let it be set1. One way to do this is to keep picking the maximum element which has not been yet picked and it to set1 and increase its sum. If at any point the remaining sum is less than or equal to the maximum element which is going to picked, add the element equal to remaining sum directly to set1, which will make the sum of set1 half of total sum. Now just add non-picked elements into set2.

Step-by-step algorithm:

• Calculate the total sum of numbers from 1 to N using: total_sum = (N * (N + 1)) / 2.
• Check if the total sum is odd. If it is, print “NO” because dividing an odd sum into two equal halves is not possible, else print “YES” as division is possible.
• Initialize an empty vector for set1 and set2, and a vector ‘vis’ to keep track of picked elements.
• Use a while loop to fill set1. Inside the loop, check if the remaining sum (total_sum / 2 – set1_sum) is greater than the maximum element not yet picked. If true, add the maximum element to set1 and update set1_sum and max_element. If false, add the remaining sum directly to set1 and set set1_sum to total_sum / 2.
• After filling set1, iterate through the numbers from 1 to n. If a number is not visited, add it to set2. and then print set1 and set2.

Below is the implementation of the algorithm:

C++

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;

void TwoSets(int n)
{
    // Calculate the total sum of numbers from 1 to n
    ll total_sum = (1ll * n * (1ll * n + 1)) / 2;

    // Check if the total sum is odd, if so, it's not
    // possible to divide into two equal halves
    if (total_sum % 2 != 0) {
        cout << "NO\n";
    }
    else {
        cout << "YES\n";

        // Initialize vectors for set1, set2, and a vector
        // to keep track of visited elements
        vector<int> set1, set2;
        vector<int> vis(n + 1, 0);

        // Initialize set1_sum and max_element
        ll set1_sum = 0;
        ll max_element = n;

        // Loop to fill set1
        while (set1_sum < total_sum / 2) {
            // Calculate remaining sum needed to make
            // set1_sum equal to total_sum / 2
            ll remaining_sum = total_sum / 2 - set1_sum;

            // If remaining_sum is greater than the maximum
            // element, add the maximum element to set1
            if (remaining_sum > max_element) {
                set1.push_back(max_element);
                vis[max_element] = 1;
                set1_sum += max_element;
                max_element--;
            }
            else {
                // If remaining_sum is less than or equal to
                // the maximum element, add remaining_sum to
                // set1
                set1.push_back(remaining_sum);
                vis[remaining_sum] = 1;
                set1_sum = total_sum / 2;
            }
        }

        // Loop to fill set2 with non-picked elements
        for (int i = 1; i <= n; i++) {
            if (vis[i] == 0) {
                set2.push_back(i);
            }
        }

        // Print the size and elements of set1
        cout << set1.size() << "\n";
        for (auto x : set1) {
            cout << x << " ";
        }
        cout << "\n";

        // Print the size and elements of set2
        cout << set2.size() << "\n";
        for (auto x : set2) {
            cout << x << " ";
        }
    }
}

// Driver Code
int main()
{
    int n=7;
    TwoSets(n);
}

Java

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class Main {

    static void TwoSets(int n)
    {
        // Calculate the total sum of numbers from 1 to n
        long totalSum = (1L * n * (1L * n + 1)) / 2;

        // Check if the total sum is odd, if so, it's not
        // possible to divide into two equal halves
        if (totalSum % 2 != 0) {
            System.out.println("NO");
        }
        else {
            System.out.println("YES");

            // Initialize lists for set1, set2, and a list
            // to keep track of visited elements
            List<Integer> set1 = new ArrayList<>();
            List<Integer> set2 = new ArrayList<>();
            boolean[] vis = new boolean[n + 1];

            // Initialize set1Sum and maxElement
            long set1Sum = 0;
            long maxElement = n;

            // Loop to fill set1
            while (set1Sum < totalSum / 2) {
                // Calculate remaining sum needed to make
                // set1Sum equal to totalSum / 2
                long remainingSum = totalSum / 2 - set1Sum;

                // If remainingSum is greater than the
                // maximum element, add the maximum element
                // to set1
                if (remainingSum > maxElement) {
                    set1.add((int)maxElement);
                    vis[(int)maxElement] = true;
                    set1Sum += maxElement;
                    maxElement--;
                }
                else {
                    // If remainingSum is less than or equal
                    // to the maximum element, add
                    // remainingSum to set1
                    set1.add((int)remainingSum);
                    vis[(int)remainingSum] = true;
                    set1Sum = totalSum / 2;
                }
            }

            // Loop to fill set2 with non-picked elements
            for (int i = 1; i <= n; i++) {
                if (!vis[i]) {
                    set2.add(i);
                }
            }

            // Print the size and elements of set1
            System.out.println(set1.size());
            for (int x : set1) {
                System.out.print(x + " ");
            }
            System.out.println();

            // Print the size and elements of set2
            System.out.println(set2.size());
            for (int x : set2) {
                System.out.print(x + " ");
            }
        }
    }

    // Driver Code
    public static void main(String[] args)
    {
        int n = 7;
        TwoSets(n);
    }
}

// This code is contributed by shivamgupta0987654321

Python3

def two_sets(n):
    # Calculate the total sum of numbers from 1 to n
    total_sum = (n * (n + 1)) // 2

    # Check if the total sum is odd, if so, it's not possible to divide into two equal halves
    if total_sum % 2 != 0:
        print("NO")
    else:
        print("YES")

        # Initialize lists for set1, set2, and a list to keep track of visited elements
        set1, set2 = [], []
        vis = [0] * (n + 1)

        # Initialize set1_sum and max_element
        set1_sum = 0
        max_element = n

        # Loop to fill set1
        while set1_sum < total_sum // 2:
            # Calculate remaining sum needed to make set1_sum equal to total_sum / 2
            remaining_sum = total_sum // 2 - set1_sum

            # If remaining_sum is greater than the maximum element, add the maximum element to set1
            if remaining_sum > max_element:
                set1.append(max_element)
                vis[max_element] = 1
                set1_sum += max_element
                max_element -= 1
            else:
                # If remaining_sum is less than or equal to the maximum element, add remaining_sum to set1
                set1.append(remaining_sum)
                vis[remaining_sum] = 1
                set1_sum = total_sum // 2

        # Loop to fill set2 with non-picked elements
        for i in range(1, n + 1):
            if vis[i] == 0:
                set2.append(i)

        # Print the size and elements of set1
        print(len(set1))
        print(" ".join(map(str, set1)))

        # Print the size and elements of set2
        print(len(set2))
        print(" ".join(map(str, set2)))


# Driver Code
if __name__ == "__main__":
    n = 7
    two_sets(n)

C#

using System;
using System.Collections.Generic;

public class Program
{
    static void TwoSets(int n)
    {
        // Calculate the total sum of numbers from 1 to n
        long totalSum = (long)n * ((long)n + 1) / 2;

        // Check if the total sum is odd, if so, it's not
        // possible to divide into two equal halves
        if (totalSum % 2 != 0)
        {
            Console.WriteLine("NO");
        }
        else
        {
            Console.WriteLine("YES");

            // Initialize lists for set1, set2, and a list
            // to keep track of visited elements
            List<int> set1 = new List<int>();
            List<int> set2 = new List<int>();
            List<int> visited = new List<int>(new int[n + 1]);

            // Initialize set1Sum and maxElement
            long set1Sum = 0;
            long maxElement = n;

            // Loop to fill set1
            while (set1Sum < totalSum / 2)
            {
                // Calculate remaining sum needed to make
                // set1Sum equal to totalSum / 2
                long remainingSum = totalSum / 2 - set1Sum;

                // If remainingSum is greater than the maximum
                // element, add the maximum element to set1
                if (remainingSum > maxElement)
                {
                    set1.Add((int)maxElement);
                    visited[(int)maxElement] = 1;
                    set1Sum += maxElement;
                    maxElement--;
                }
                else
                {
                    // If remainingSum is less than or equal to
                    // the maximum element, add remainingSum to
                    // set1
                    set1.Add((int)remainingSum);
                    visited[(int)remainingSum] = 1;
                    set1Sum = totalSum / 2;
                }
            }

            // Loop to fill set2 with non-picked elements
            for (int i = 1; i <= n; i++)
            {
                if (visited[i] == 0)
                {
                    set2.Add(i);
                }
            }

            // Print the size and elements of set1
            Console.WriteLine(set1.Count);
            Console.WriteLine(string.Join(" ", set1));

            // Print the size and elements of set2
            Console.WriteLine(set2.Count);
            Console.WriteLine(string.Join(" ", set2));
        }
    }

    // Driver Code
    public static void Main()
    {
        int n = 7;
        TwoSets(n);
    }
}

JavaScript

function twoSets(n) {
    // Calculate the total sum of numbers from 1 to n
    const totalSum = (n * (n + 1)) / 2;

    // Check if the total sum is odd, if so, it's not possible to divide into two equal halves
    if (totalSum % 2 !== 0) {
        console.log("NO");
    } else {
        console.log("YES");

        // Initialize arrays for set1, set2, and a boolean array to keep track of visited elements
        const set1 = [];
        const set2 = [];
        const vis = new Array(n + 1).fill(0);

        // Initialize set1Sum and maxElement
        let set1Sum = 0;
        let maxElement = n;

        // Loop to fill set1
        while (set1Sum < totalSum / 2) {
            // Calculate remaining sum needed to make set1Sum equal to totalSum / 2
            const remainingSum = totalSum / 2 - set1Sum;

            // If remainingSum is greater than the maximum element, add the maximum element to set1
            if (remainingSum > maxElement) {
                set1.push(maxElement);
                vis[maxElement] = 1;
                set1Sum += maxElement;
                maxElement -= 1;
            } else {
                // If remainingSum is less than or equal to the maximum element, add remainingSum to set1
                set1.push(remainingSum);
                vis[remainingSum] = 1;
                set1Sum = totalSum / 2;
            }
        }

        // Loop to fill set2 with non-picked elements
        for (let i = 1; i <= n; i++) {
            if (vis[i] === 0) {
                set2.push(i);
            }
        }

        // Print the size and elements of set1
        console.log(set1.length);
        console.log(set1.join(" "));

        // Print the size and elements of set2
        console.log(set2.length);
        console.log(set2.join(" "));
    }
}

// Driver Code
const n = 7;
twoSets(n);

Output

YES
3
7 6 1 
4
2 3 4 5 

Time Complexity: O(N), where N is the total number of integers given in the input.
Auxiliary Space: O(N)

Please refer to the complete guide of CSES Problem Set Solutions.