leetcode 565. 数组嵌套（C++、python）

索引从0开始长度为N的数组A，包含0到N - 1的所有整数。找到并返回最大的集合S，S[i] = {A[i], A[A[i]], A[A[A[i]]], ... }且遵守以下的规则。

假设选择索引为i的元素A[i]为S的第一个元素，S的下一个元素应该是A[A[i]]，之后是A[A[A[i]]]... 以此类推，不断添加直到S出现重复的元素。

示例 1:

输入: A = [5,4,0,3,1,6,2]
输出: 4
解释: 
A[0] = 5, A[1] = 4, A[2] = 0, A[3] = 3, A[4] = 1, A[5] = 6, A[6] = 2.

其中一种最长的 S[K]:
S[0] = {A[0], A[5], A[6], A[2]} = {5, 6, 2, 0}

注意:

N是[1, 20,000]之间的整数。

A中不含有重复的元素。

A中的元素大小在[0, N-1]之间。

思路：遍历元素时（假设为火车头），用一个映射关系记录从火车头开始最长的集合数量，如果该元素不是集合的火车头，那么直接记为-1，继续处理后一个元素；如果从火车头开始，到某个位置时，该位置映射的值大于0，那么直接将火车头累积的值加上该位置映射的值，就得到火车头的映射。

C++

class Solution {
public:
    int arrayNesting(vector<int>& nums) 
    {
        int n=nums.size();
        int ans=0;
        map<int,int> vec;
        for(int i=0;i<n;i++)
        {
            if(vec[nums[i]])
            {
                continue;
            }
            else
            {
                int res=1;
                int start=nums[i];
                int tmp=nums[start];
                while(tmp!=start)
                {
                    if(vec[tmp]>0)
                    {
                        res+=vec[tmp];
                        vec[start]=res;
                        ans=max(res,ans);
                        break;
                    }
                    else
                    {
                        vec[tmp]=-1;
                        res++;
                        tmp=nums[tmp];
                    }
                }
                ans=max(res,ans);
                vec[start]=res;
            }
        }
        return ans;
    }
};

python

class Solution:
    def arrayNesting(self, nums: List[int]) -> int:
        n=len(nums)
        dic={}
        ans=0
        for i in range(n):
            start=nums[i]
            if start in dic:
                continue
            else:
                res=1
                tmp=nums[start]
                while tmp!=start:
                    if tmp in dic:
                        res+=dic[tmp]
                        dic[start]=res
                        ans=max(ans,res)
                        break
                    else:
                        dic[tmp]=-1
                        res+=1
                        tmp=nums[tmp]
                ans=max(ans,res)
                dic[start]=res
        return ans