LeetCode in Kotlin

2076. Process Restricted Friend Requests

Hard

You are given an integer n indicating the number of people in a network. Each person is labeled from 0 to n - 1.

You are also given a 0-indexed 2D integer array restrictions, where restrictions[i] = [xi, yi] means that person xi and person yi cannot become friends, either directly or indirectly through other people.

Initially, no one is friends with each other. You are given a list of friend requests as a 0-indexed 2D integer array requests, where requests[j] = [uj, vj] is a friend request between person uj and person vj.

A friend request is successful if uj and vj can be friends. Each friend request is processed in the given order (i.e., requests[j] occurs before requests[j + 1]), and upon a successful request, uj and vj become direct friends for all future friend requests.

Return a boolean array result, where each result[j] is true if the jth friend request is successful or false if it is not.

Note: If uj and vj are already direct friends, the request is still successful.

Example 1:

Input: n = 3, restrictions = [[0,1]], requests = [[0,2],[2,1]]

Output: [true,false]

Explanation:

Request 0: Person 0 and person 2 can be friends, so they become direct friends.

Request 1: Person 2 and person 1 cannot be friends since person 0 and person 1 would be indirect friends (1–2–0).

Example 2:

Input: n = 3, restrictions = [[0,1]], requests = [[1,2],[0,2]]

Output: [true,false]

Explanation:

Request 0: Person 1 and person 2 can be friends, so they become direct friends.

Request 1: Person 0 and person 2 cannot be friends since person 0 and person 1 would be indirect friends (0–2–1).

Example 3:

Input: n = 5, restrictions = [[0,1],[1,2],[2,3]], requests = [[0,4],[1,2],[3,1],[3,4]]

Output: [true,false,true,false]

Explanation:

Request 0: Person 0 and person 4 can be friends, so they become direct friends.

Request 1: Person 1 and person 2 cannot be friends since they are directly restricted.

Request 2: Person 3 and person 1 can be friends, so they become direct friends.

Request 3: Person 3 and person 4 cannot be friends since person 0 and person 1 would be indirect friends (0–4–3–1).

Constraints:

Solution

@Suppress("NAME_SHADOWING")
class Solution {
    fun friendRequests(n: Int, restrictions: Array<IntArray>, requests: Array<IntArray>): BooleanArray {
        // Check for each request whether it can cause conflict or not
        val uf = UnionFind(n)
        val res = BooleanArray(requests.size)
        for (i in requests.indices) {
            val p1 = uf.findParent(requests[i][0])
            val p2 = uf.findParent(requests[i][1])
            if (p1 == p2) {
                res[i] = true
                continue
            }
            // Check whether the current request will violate any restriction or not
            var flag = true
            for (restrict in restrictions) {
                val r1 = uf.findParent(restrict[0])
                val r2 = uf.findParent(restrict[1])
                if (r1 == p1 && r2 == p2 || r1 == p2 && r2 == p1) {
                    flag = false
                    break
                }
            }
            if (flag) {
                res[i] = true
                // Union
                uf.parent[p1] = p2
            }
        }
        return res
    }

    private class UnionFind(n: Int) {
        var parent: IntArray = IntArray(n)

        init {
            for (i in 0 until n) {
                parent[i] = i
            }
        }

        fun findParent(user: Int): Int {
            var user = user
            while (parent[user] != user) {
                parent[user] = parent[parent[user]]
                user = parent[user]
            }
            return user
        }
    }
}