Find the tasks completed by soldiers based on their ranks

Given an array ranks[] of size N denoting the ranks of soldiers and tasks[] of size M where tasks[i] represent the amount of time (in seconds) required to complete i^th task. Every second, a task enters into the system. Assigning tasks to soldiers depends on the rank of soldiers, the task will be assigned to a free soldier with the smallest rank, and in the case of a tie, it is assigned to the free soldier with the smallest index. The task is to return the array of size M which will represent the index of soldiers to which a particular task will be assigned.

Note: If all soldiers are busy then the task has to wait for a soldier to get free from their current task.

Examples:

Input: n = 3, m = 4, ranks[] = {3, 1, 2}, tasks[] = {1, 3, 2, 1}
Output: {1, 1, 2, 0}
Explanation: The output array is as follows:

• tasks[0] = 1, entering the system at 0 seconds. Assigning to smallest rank from ranks i.e. at index 1.
• tasks[1] = 3, entering the system at 1 second. Assigning to smallest rank from ranks i.e. at index 1 as it was busy for only 1 second.
• tasks[2] = 2, entering the system at 2 seconds. Assigning to smallest rank from ranks i.e. at index 2 (not index 1 because it is busy performing tasks[1])
• tasks[3] = 1, entering the system at 3 seconds. Assigning to smallest rank from ranks i.e. at index 0.

Input: n = 4, m = 4, ranks[] = {3, 1, 2, 4}, tasks[] = {1, 1, 2, 1}
Output: {1, 1, 1, 2}
Explanation: The output array is as follows:

• tasks[0] = 1, entering the system at 0 seconds. Assigning to smallest rank from ranks i.e. at index 1.
• tasks[1] = 1, entering the system at 1 second. Assigning to smallest rank from ranks i.e. at index 1 as it was busy for only 1 second.
• tasks[2] = 2, entering the system at 2 seconds. Assigning to smallest rank from ranks i.e. at index 1 as it was bus for only 1 second.
• tasks[3] = 1, entering the system at 3 seconds. Assigning to smallest rank from ranks i.e. at index 2 (not index 1 because it is busy performing tasks[2]).

Approach: This can be solved with the following approach:

Maintain a min-heap (priority queue) pq to store pairs of the ranks of soldiers along with their indices and an ans array to store the indices of soldiers. This min heap will have the top soldier with the smallest rank as well as index. Also maintain a map free which stores when particular soldier would be getting free. We will keep track of the time by using a timer variable. Now iterate over every task and see if we have any soldier who is free at the time by checking the min heap.

• If yes, then pop the soldier from the min heap and update the answer. Append the popped soldier to the free map at index (timer + time required to complete the current task). This will tell us when will this soldier get free.
• If no, that is the min heap is empty. Then, it means that there are no soldiers who are available to perform the task, so we increment the timer to the time when a soldier gets free.

We will be freeing the soldiers after they complete their tasks using the free map. Keep on assigning the tasks based on the above steps and keep updating the ans[] array. Finally, return the ans[] array.

Steps involved in the implementation of the code:

• Iterate over ranks[] and store pair {ranks[i], i} in priority queue pq (min heap).
• Initialize a timer variable to keep track of the current time.
• For each tasks[i]:
  • Check the free map to see if there are any soldiers who are getting free at timer seconds. If yes, then pop those soldiers from the free map and add them to pq.
  • Check pq if there are soldiers available to perform the current task
    • If pq is empty, then increment the timer to the time when a soldiers gets free and continue doing the current task.
    • Else, pop the top soldier from pq and add it to the free map at index (timer + tasks[i]), update the ans[] array, increment the timer by one and move to the next task.
• Return the ans[] array.

Below is the implementation of the code:

C++

// C++ code for the above approach

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

vector<int> assign_tasks(vector<int> ranks,
                         vector<int> tasks)
{

    // min-heap to store the soldiers based on
    // their ranks and indices
    priority_queue<pair<int, int>, vector<pair<int, int> >,
                   greater<pair<int, int> > >
        pq;

    // Insert all the free soldiers
    // to the min heap
    for (int i = 0; i < ranks.size(); i++)
        pq.push({ ranks[i], i });

    vector<int> ans;

    // map to store soldiers who are busy
    // free[i] will consists of all soldiers
    // who will get free at ith second
    map<int, vector<pair<int, int> > > free;
    int timer = 0;
    for (int i = 0; i < tasks.size(); i++) {

        // Insert the soldiers who were supposed
        // to be free at time = timer seconds
        for (auto busy_soldier : free[timer])
            pq.push(busy_soldier);
        free.erase(timer);

        // If no soldier is free, then wait
        // till a soldier gets free
        if (pq.empty()) {
            timer = free.begin()->first;
            i--;
        }
        else {

            // If a soldier is free, then allot
            // the task to the soldier with
            // smallest rank and smallest index
            int soldier_rank = pq.top().first;
            int soldier_index = pq.top().second;
            pq.pop();
            ans.push_back(soldier_index);
            free[timer + tasks[i]].push_back(
                { soldier_rank, soldier_index });
            timer += 1;
        }
    }
    return ans;
}

// Drivers code
int main()
{
    int n = 3;
    int m = 4;
    vector<int> ranks = { 3, 1, 2 };
    vector<int> tasks = { 1, 3, 2, 1 };

    // Functional call
    vector<int> ans = assign_tasks(ranks, tasks);

    for (int i = 0; i < m; i++)
        cout << ans[i] << " ";
    return 0;
}

Java

import java.util.*;

class Main {
    public static List<Integer> assignTasks(List<Integer> ranks, List<Integer> tasks) {
        // Create a priority queue to store soldiers based on their ranks
        // Priority is determined first by rank and then by their original index.
        PriorityQueue<AbstractMap.SimpleEntry<Integer, Integer>> pq = new PriorityQueue<>(
            (a, b) -> a.getKey() != b.getKey() ? a.getKey() - b.getKey() : a.getValue() - b.getValue()
        );

        // Initialize the priority queue with soldiers and their ranks.
        for (int i = 0; i < ranks.size(); i++) {
            pq.offer(new AbstractMap.SimpleEntry<>(ranks.get(i), i));
        }

        List<Integer> ans = new ArrayList<>(); // Initialize the list to store the assigned tasks.
        Map<Integer, List<AbstractMap.SimpleEntry<Integer, Integer>>> free = new HashMap<>(); // Initialize a map to track free soldiers.
        int timer = 0; // Initialize a timer to keep track of time.

        for (int i = 0; i < tasks.size(); i++) {
            // Release soldiers who have completed their tasks at the current time.
            for (AbstractMap.SimpleEntry<Integer, Integer> busySoldier : free.getOrDefault(timer, new ArrayList<>())) {
                pq.offer(busySoldier);
            }
            free.remove(timer);

            if (pq.isEmpty()) {
                // If no soldier is available, find the earliest time a soldier becomes free and adjust the loop accordingly.
                timer = Objects.requireNonNull(free.keySet().stream().min(Integer::compareTo).orElse(null));
                i--;
            } else {
                // Assign a task to the next available soldier from the priority queue.
                AbstractMap.SimpleEntry<Integer, Integer> soldier = pq.poll();
                ans.add(soldier.getValue());
                // Update the free time for the soldier based on the task duration.
                free.computeIfAbsent(timer + tasks.get(i), k -> new ArrayList<>()).add(soldier);
                timer += 1;
            }
        }

        return ans;
    }

    public static void main(String[] args) {
        int n = 3;
        int m = 4;
        List<Integer> ranks = Arrays.asList(3, 1, 2);
        List<Integer> tasks = Arrays.asList(1, 3, 2, 1);

        List<Integer> ans = assignTasks(ranks, tasks);

        // Print the assigned soldiers for each task.
        for (int i = 0; i < m; i++) {
            System.out.print(ans.get(i) + " ");
        }
    }
}

Python3

import heapq

def assign_tasks(ranks, tasks):
    # Min-heap to store the soldiers based on their ranks and indices
    pq = []

    # Insert all the free soldiers to the min heap
    for i in range(len(ranks)):
        heapq.heappush(pq, (ranks[i], i))

    ans = []

    # Dictionary to store soldiers who are busy
    # free[i] will consist of all soldiers who will get free at ith second
    free = {}
    timer = 0

    for i in range(len(tasks)):
        # Insert the soldiers who were supposed to be free at time = timer seconds
        if timer in free:
            for busy_soldier in free[timer]:
                heapq.heappush(pq, busy_soldier)
            del free[timer]

        # If no soldier is free, then wait till a soldier gets free
        if not pq:
            timer = min(free.keys())
            i -= 1
        else:
            # If a soldier is free, then allot the task to the soldier with the smallest rank and smallest index
            soldier_rank, soldier_index = heapq.heappop(pq)
            ans.append(soldier_index)
            if timer + tasks[i] in free:
                free[timer + tasks[i]].append((soldier_rank, soldier_index))
            else:
                free[timer + tasks[i]] = [(soldier_rank, soldier_index)]
            timer += 1

    return ans

# Driver code
n = 3
m = 4
ranks = [3, 1, 2]
tasks = [1, 3, 2, 1]

# Function call
ans = assign_tasks(ranks, tasks)

for i in range(m):
    print(ans[i], end=" ")

C#

using System;
using System.Collections.Generic;
using System.Linq;

class Program {
    static List<int> AssignTasks(List<int> ranks,
                                 List<int> tasks, int n)
    {
        // Min-heap to store the soldiers based on their
        // ranks and indices
        var pq = new SortedDictionary<int, int>();

        // Insert all the free soldiers to the min heap
        for (int i = 0; i < ranks.Count; i++) {
            pq[ranks[i]] = i;
        }

        var ans = new List<int>();

        // Dictionary to store soldiers who are busy
        // free[i] will consist of all soldiers who will get
        // free at ith second
        var free
            = new Dictionary<int,
                             List<Tuple<int, int> > >();
        int timer = 0;
        for (int i = 0; i < tasks.Count; i++) {
            // Insert the soldiers who were supposed to be
            // free at time = timer seconds
            if (free.ContainsKey(timer)) {
                foreach(var busySoldier in free[timer])
                {
                    pq[busySoldier.Item1]
                        = busySoldier.Item2;
                }
                free.Remove(timer);
            }

            // If no soldier is free, then wait till a
            // soldier gets free
            if (pq.Count == 0) {
                timer = free.Keys.Min();
                i--;
            }
            else {
                // If a soldier is free, then allot the task
                // to the soldier with smallest rank and
                // smallest index
                int soldierRank = pq.Keys.Min();
                int soldierIndex = pq[soldierRank];
                pq.Remove(soldierRank);

                ans.Add(soldierIndex);

                if (!free.ContainsKey(timer + tasks[i])) {
                    free[timer + tasks[i]]
                        = new List<Tuple<int, int> >();
                }

                free[timer + tasks[i]].Add(
                    new Tuple<int, int>(soldierRank,
                                        soldierIndex));
                timer++;
            }
        }
        return ans;
    }

    static void Main()
    {
        int n = 3;
        int m = 4;
        List<int> ranks = new List<int>{ 3, 1, 2 };
        List<int> tasks = new List<int>{ 1, 3, 2, 1 };

        // Functional call
        List<int> ans = AssignTasks(ranks, tasks, n);

        for (int i = 0; i < m; i++) {
            Console.Write(ans[i] + " ");
        }
    }
}

Javascript

// Implementation of a Priority Queue class
class PriorityQueue {
    constructor() {
        this.data = [];
    }

    // Add an element with a specified priority to the queue
    enqueue(element, priority) {
        this.data.push({ element, priority });
        this.sort();
    }

    // Remove and return the element with the highest priority
    dequeue() {
        return this.data.shift().element;
    }

    // Sort the queue based on priority
    sort() {
        this.data.sort((a, b) => a.priority - b.priority);
    }

    // Check if the queue is empty
    isEmpty() {
        return this.data.length === 0;
    }
}

// Function to assign tasks to soldiers based on their ranks and tasks' durations
function assignTasks(ranks, tasks) {
    // Create a Priority Queue to store soldiers based on their ranks
    const pq = new PriorityQueue();

    // Initialize the Priority Queue with soldiers and their ranks
    for (let i = 0; i < ranks.length; i++) {
        pq.enqueue(i, ranks[i]);
    }

    // Array to store the assigned tasks to soldiers
    const ans = [];

    // Dictionary to store soldiers who are busy
    const free = {};

    // Variable to keep track of time
    let timer = 0;

    // Iterate through each task
    for (let i = 0; i < tasks.length; i++) {
        // Check if there are soldiers who became free at the current time
        if (free[timer]) {
            // Enqueue the free soldiers back to the Priority Queue
            free[timer].forEach((busySoldier) => {
                pq.enqueue(busySoldier.index, busySoldier.rank);
            });
            // Remove the entry from the dictionary
            delete free[timer];
        }

        // If no soldiers are free, wait until a soldier becomes free
        if (pq.isEmpty()) {
            // Find the minimum time when a soldier becomes free
            timer = Math.min(...Object.keys(free));
            i--; // Decrement i to reprocess the current task with the newly available soldiers
        } else {
            // Assign the task to the soldier with the highest priority (smallest rank)
            const soldierIndex = pq.dequeue();
            ans.push(soldierIndex);

            // Calculate the end time of the current task
            const taskEndTime = timer + tasks[i];

            // Initialize the entry in the free dictionary if not exists
            if (!free[taskEndTime]) {
                free[taskEndTime] = [];
            }

            // Add the soldier to the list of soldiers who will be free at the task end time
            free[taskEndTime].push({ rank: ranks[soldierIndex], index: soldierIndex });
            timer++; // Move to the next second
        }
    }

    return ans;
}

// Example usage
const ranks = [3, 1, 2];
const tasks = [1, 3, 2, 1];

// Assign tasks and get the result
const ans = assignTasks(ranks, tasks);

// Display the result
console.log(ans.join(" "));

Output

1 1 2 0 

Time Complexity: O(N log N), where N is the size of ranks[] array
Auxiliary Space: O(M + N), where M is the size of tasks[] array